DuckDB PARTITION_BY 完整指南：一条 SQL 搞定数据分区，查询加速 50 倍

-- 你的查询：只看 2026 年 5 月的数据
SELECT
    product_category,
    COUNT(*) AS order_count,
    SUM(amount) AS total_revenue
FROM orders
WHERE order_date >= '2026-05-01' AND order_date < '2026-06-01'
GROUP BY product_category
ORDER BY total_revenue DESC;

-- 假设你有一个原始 orders 表
CREATE TABLE raw_orders AS
SELECT * FROM read_parquet('orders_raw.parquet');

-- 一条 SQL 把全表按 order_date 列分区写入
COPY (
    SELECT * FROM raw_orders
) TO 'orders_partitioned/'
(FORMAT PARQUET, PARTITION_BY (order_date));

COPY (
    SELECT * FROM raw_orders
) TO 'orders_partitioned/'
(FORMAT PARQUET, 
 PARTITION_BY (order_date),
 ROWS_PER_GROUP 100000);

COPY (
    SELECT * FROM raw_orders
) TO 'orders_partitioned/'
(FORMAT PARQUET, 
 PARTITION_BY (order_date, region));

SELECT
    product_category,
    COUNT(*) AS order_count,
    SUM(amount) AS total_revenue
FROM read_parquet('orders_partitioned/order_date=2026-05/region=华东/*.parquet')
GROUP BY product_category;

-- 每月初跑一次增量导入
COPY (
    SELECT * FROM staging_new_orders
    WHERE order_date >= '2026-06-01'
) TO 'orders_partitioned/'
(FORMAT PARQUET, PARTITION_BY (order_date), OVERWRITE_MODE 'IGNORE');

import duckdb
import time

con = duckdb.connect(":memory:")

# 场景 A：非分区查询 —— 扫描全部数据
start = time.time()
result_a = con.execute("""
    SELECT product_category, COUNT(*), SUM(amount)
    FROM read_parquet('orders_raw.parquet')
    WHERE order_date >= '2026-05-01' AND order_date < '2026-06-01'
    GROUP BY product_category
""").fetchdf()
time_a = time.time() - start

# 场景 B：分区查询 —— 自动裁剪
start = time.time()
result_b = con.execute("""
    SELECT product_category, COUNT(*), SUM(amount)
    FROM read_parquet('orders_partitioned/*.parquet')
    WHERE order_date >= '2026-05-01' AND order_date < '2026-06-01'
    GROUP BY product_category
""").fetchdf()
time_b = time.time() - start

print(f"非分区查询：{time_a:.2f} 秒")
print(f"分区裁剪查询：{time_b:.2f} 秒")
print(f"加速比：{time_a/time_b:.1f}x")

COPY (SELECT * FROM raw_orders) TO 'orders_partitioned/'
(FORMAT PARQUET, PARTITION_BY (order_date), ROWS_PER_GROUP 100000);

-- ✅ 有效裁剪：分区列名精确匹配
WHERE order_date = '2026-05-01'

-- ⚠️ 可能不裁剪：分区列在函数中被包裹
WHERE EXTRACT(YEAR FROM order_date) = 2026

-- ✅ 正确：匹配分区目录下的文件
FROM read_parquet('orders_partitioned/order_date=2026-05/*.parquet')

-- ❌ 错误：不会触发分区裁剪
FROM read_parquet('orders_partitioned/*.parquet')
WHERE order_date = '2026-05-01'
-- 虽然结果正确，但会扫描所有分区目录

import duckdb
from pathlib import Path
import time

con = duckdb.connect(":memory:")

# ========== 第一步：生成测试数据 ==========
print("Step 1: 生成测试数据...")
con.execute("""
    CREATE TABLE raw_orders AS
    SELECT
        generate_series AS order_id,
        DATE '2020-01-01' + (random() * 2190)::INTEGER AS order_date,
        ('电子产品' || (random()*5)::INTEGER) AS product_category,
        (random() * 1000 + 10)::DECIMAL(10,2) AS amount,
        CASE 
            WHEN random() < 0.3 THEN '华东'
            WHEN random() < 0.6 THEN '华南'
            ELSE '华北'
        END AS region
    FROM generate_series(1, 5000000)
""")
print(f"生成 {con.execute('SELECT count(*) FROM raw_orders').fetchone()[0]} 行数据")

# ========== 第二步：创建分区目录 ==========
print("Step 2: 转成分区格式...")
part_dir = Path("orders_partitioned")
if part_dir.exists():
    import shutil
    shutil.rmtree(part_dir)

con.execute("""
    COPY (SELECT * FROM raw_orders) TO 'orders_partitioned/'
    (FORMAT PARQUET, PARTITION_BY (order_date));
""")
print(f"分区目录结构：")
for d in sorted(part_dir.iterdir()):
    if d.is_dir():
        file_count = len(list(d.glob("*.parquet")))
        print(f"  {d.name}/ -> {file_count} 个文件")

# ========== 第三步：对比查询性能 ==========
print("\nStep 3: 性能对比...")

# 非分区查询
start = time.time()
result_a = con.execute("""
    SELECT product_category, COUNT(*), SUM(amount)
    FROM read_parquet('orders_raw_single.parquet')
    WHERE order_date >= '2026-05-01' AND order_date < '2026-06-01'
    GROUP BY product_category
""").fetchdf()
time_a = time.time() - start

# 分区查询
start = time.time()
result_b = con.execute("""
    SELECT product_category, COUNT(*), SUM(amount)
    FROM read_parquet('orders_partitioned/order_date=2026-05/*.parquet')
    GROUP BY product_category
""").fetchdf()
time_b = time.time() - start

print(f"非分区查询：{time_a:.2f} 秒")
print(f"分区裁剪查询：{time_b:.2f} 秒")