import duckdb

# 最简单的方式：一行代码导入
con = duckdb.connect("big_data.duckdb")
con.execute("CREATE TABLE events AS SELECT * FROM read_csv_auto('events_2026.csv')")

# 读取目录下所有 CSV 文件（自动并行）
con.execute("""
    CREATE TABLE events AS
    SELECT * FROM read_csv_auto('events_2026/*.csv')
""")

# 按时间范围筛选后写入（谓词下推，只读需要的数据）
con.execute("""
    CREATE TABLE events_2026_q3 AS
    SELECT * FROM read_csv_auto('events_2026/*.csv')
    WHERE event_date >= '2026-07-01'
      AND event_date < '2026-10-01'
""")

# 探测前 1000 行的 schema
schema = con.execute("""
    SELECT column_name, column_type
    FROM (SELECT * FROM read_csv_auto('events_2026/*.csv', SAMPLE_SIZE=1000))
    LIMIT 0
""").fetchdf()

print(schema)
#   column_name    column_type
# 0  event_id       BIGINT
# 1  user_id        BIGINT
# 2  event_type     VARCHAR
# 3  event_date     DATE
# 4  amount         DOUBLE
# ...

# 使用检测到的 schema 精确导入
con.execute("""
    CREATE TABLE events (
        event_id BIGINT,
        user_id BIGINT,
        event_type VARCHAR,
        event_date DATE,
        amount DOUBLE
    )
""")
con.execute("""
    INSERT INTO events
    SELECT * FROM read_csv_auto('events_2026/*.csv', TYPES(...))
""")

# 直接读取 gzip 压缩的 CSV（自动解压 + 并行解析）
con.execute("CREATE TABLE events AS SELECT * FROM read_csv_auto('events_2026.csv.gz')")

# 直接读取 gzip 压缩的 JSON
con.execute("CREATE TABLE events AS SELECT * FROM read_json_auto('events_2026.json.gz')")

con = duckdb.connect("big_data.duckdb", config={
    'threads': '8',              # 使用全部 8 个核心
    'max_memory': '24GB',        # 限制内存使用上限
    'temp_directory': '/tmp/duckdb_temp',  # 溢出到磁盘的目录
    'parallel': '8'              # 并行读取线程数
})

# 并行读取多个 gzip 文件
con.execute("""
    CREATE TABLE events AS
    SELECT * FROM read_csv_auto('events_2026/*.csv.gz', parallel=8)
""")

import duckdb

con = duckdb.connect("big_data.duckdb")

# 创建表结构
con.execute("""
    CREATE TABLE events (
        event_id BIGINT,
        user_id BIGINT,
        event_type VARCHAR,
        event_date DATE,
        amount DOUBLE,
        metadata JSON
    )
""")

# 使用 Appender 批量写入（比 INSERT 快 20x）
appender = con.appender("events")

# 分批写入（每批 100 万行）
batch_size = 1_000_000
for i in range(1000):  # 1000 批 = 10 亿行
    # 模拟数据生成（实际场景中替换为你的数据源）
    batch = generate_batch(batch_size)  # 返回 DataFrame 或 list of tuples
    appender.append_df(batch)
    if i % 50 == 0:
        print(f"Progress: {i*batch_size:,} / 1,000,000,000")

appender.flush()
con.close()
print("Load complete!")

# Step 1: 读取并写入 Parquet（利用列式压缩）
con1 = duckdb.connect(":memory:")
con1.execute("""
    CREATE TABLE events AS
    SELECT * FROM read_csv_auto('events_2026/*.csv.gz')
""")
con1.execute("COPY events TO 'events_2026.parquet' (FORMAT PARQUET)")
con1.close()

# Step 2: 从 Parquet 加载到 DuckDB（速度极快）
con2 = duckdb.connect("big_data.duckdb")
con2.execute("CREATE TABLE events AS SELECT * FROM 'events_2026.parquet'")
con2.close()

#!/usr/bin/env python3
"""DuckDB 十亿级数据加载脚本"""
import duckdb
import os
import time

# 配置
DATABASE = "/data/warehouse/events.duckdb"
INPUT_PATTERN = "/data/raw/events_2026/*.csv.gz"
TEMP_DIR = "/tmp/duckdb_temp"
THREADS = 8
MAX_MEMORY = "24GB"

os.makedirs(TEMP_DIR, exist_ok=True)

# 连接配置
config = {
    'threads': str(THREADS),
    'max_memory': MAX_MEMORY,
    'temp_directory': TEMP_DIR,
    'preserve_insert_order': 'false',  # 不保证顺序，加速写入
    'allow_unsigned_extensions': 'true',
}

print(f"[*] Starting DuckDB load at {time.strftime('%Y-%m-%d %H:%M:%S')}")
print(f"[*] Config: threads={THREADS}, max_memory={MAX_MEMORY}")

con = duckdb.connect(DATABASE, config=config)

# 创建表结构
con.execute("""
    CREATE TABLE IF NOT EXISTS events (
        event_id BIGINT,
        user_id BIGINT,
        event_type VARCHAR,
        event_date DATE,
        amount DOUBLE,
        country VARCHAR,
        device VARCHAR,
        browser VARCHAR,
        os VARCHAR,
        session_id VARCHAR,
        created_at TIMESTAMP,
        metadata JSON
    )
""")

# 并行读取 gzip CSV
start = time.time()
con.execute(f"""
    CREATE TABLE events_raw AS
    SELECT * FROM read_csv_auto('{INPUT_PATTERN}',
        parallel={THREADS},
        all_varchar=false,
        auto_detect=true
    )
""")
elapsed = time.time() - start
print(f"[*] Load completed in {elapsed:.1f}s")

# 清理临时数据
con.execute("DROP TABLE events_raw")
con.execute("VACUUM")

# 验证
row_count = con.execute("SELECT COUNT(*) FROM events").fetchone()[0]
print(f"[*] Total rows: {row_count:,}")
print(f"[*] Database size: {os.path.getsize(DATABASE) / 1024**3:.2f} GB")
print(f"[*] Completed at {time.strftime('%Y-%m-%d %H:%M:%S')}")

con.close()

# 监控加载进度
con = duckdb.connect("big_data.duckdb")

# 查看当前查询状态
result = con.execute("SELECT * FROM duckdb_transactions()").fetchall()
for row in result:
    print(row)

# 监控内存使用
memory_info = con.execute("""
    SELECT * FROM duckdb_memory()
""").fetchall()
for row in memory_info:
    print(row)

# 查看执行计划（确认谓词下推）
con.execute("EXPLAIN ANALYZE SELECT * FROM events WHERE event_date >= '2026-07-01'")