pip install duckdb pandas polars

import duckdb
import pandas as pd
import polars as pl

# 初始化连接（默认使用内存数据库）
con = duckdb.connect(":memory:")

# 读取 CSV 到 DuckDB（零拷贝优势）
con.sql("""
    CREATE TABLE orders AS
    SELECT * FROM read_csv_auto('orders.csv')
""")

# 在 DuckDB 中完成过滤和聚合（SQL 更高效）
agg_result = con.sql("""
    SELECT 
        customer_id,
        COUNT(*) AS order_count,
        SUM(amount) AS total_spent,
        AVG(amount) AS avg_order_value
    FROM orders
    WHERE order_date >= '2024-01-01'
      AND status = 'completed'
    GROUP BY customer_id
""").df()  # 一行代码转为 Pandas DataFrame

print(agg_result.head())

# DuckDB 执行复杂窗口函数查询
results = con.sql("""
    WITH ranked_orders AS (
        SELECT 
            customer_id,
            order_date,
            amount,
            ROW_NUMBER() OVER (
                PARTITION BY customer_id 
                ORDER BY order_date DESC
            ) AS rn,
            LAG(amount, 1) OVER (
                PARTITION BY customer_id 
                ORDER BY order_date
            ) AS prev_amount
        FROM orders
        WHERE order_date >= '2024-01-01'
    )
    SELECT customer_id, order_date, amount, prev_amount
    FROM ranked_orders
    WHERE rn <= 3
""")

# 零拷贝转为 Polars DataFrame
df_pl = results.pl()  # 返回 Polars LazyFrame

df_eager = df_pl.collect()  # 触发执行
print(df_eager.head(6))

import time
import pyarrow as pa

# 准备 500万行测试数据
print("生成测试数据...")
test_df = pd.DataFrame({
    'id': range(5_000_000),
    'value': pd.np.random.randn(5_000_000),
    'category': pd.np.random.choice(['A', 'B', 'C'], 5_000_000),
})
con.sql("CREATE TABLE test_data AS SELECT * FROM test_df")

# 测试1: DuckDB → Pandas (df())
t0 = time.time()
result_pd = con.sql("SELECT * FROM test_data WHERE value > 0").df()
t1 = time.time()
print(f"DuckDB→Pandas: {t1-t0:.3f}s, shape={result_pd.shape}")

# 测试2: DuckDB → Polars (pl())
t0 = time.time()
result_pl = con.sql("SELECT * FROM test_data WHERE value > 0").pl().collect()
t1 = time.time()
print(f"DuckDB→Polars: {t1-t0:.3f}s, shape={result_pl.shape}")

# 测试3: Pandas → DuckDB (INSERT)
t0 = time.time()
con.sql("INSERT INTO test_data SELECT * FROM test_df")
t1 = time.time()
print(f"Pandas→DuckDB: {t1-t0:.3f}s")

# 测试4: Polars → DuckDB
t0 = time.time()
con.sql("INSERT INTO test_data SELECT * FROM result_pl")
t1 = time.time()
print(f"Polars→DuckDB:  {t1-t0:.3f}s")

# 第一步：用 SQL 快速探索数据分布
distribution = con.sql("""
    SELECT 
        category,
        COUNT(*) AS cnt,
        PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY value) AS median,
        PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY value) AS p95
    FROM test_data
    GROUP BY category
""").df()
print(distribution)

# 第二步：基于 SQL 结果，用 Pandas 进行精细计算
import numpy as np
for cat in distribution['category']:
    subset = con.sql(f"SELECT value FROM test_data WHERE category = '{cat}'").df()
    # 在 Pandas 中做自定义特征工程
    subset['log_value'] = np.log1p(subset['value'].abs())
    subset['is_outlier'] = subset['value'] > subset['value'].quantile(0.99)
    print(f"  {cat}: outliers={subset['is_outlier'].sum()}, mean_log={subset['log_value'].mean():.4f}")

# 创建 Polars DataFrame
pl_df = pl.DataFrame({
    'product_id': [101, 102, 103, 104, 105],
    'price': [29.99, 49.99, 19.99, 99.99, 39.99],
    'stock': [150, 80, 200, 30, 120],
    'category': ['Electronics', 'Home', 'Books', 'Electronics', 'Fashion'],
})

# 注册为 DuckDB 表
con.register('products', pl_df)

# 用 SQL 查询 Polars 数据
result = con.sql("""
    SELECT 
        category,
        AVG(price) AS avg_price,
        SUM(stock) AS total_stock
    FROM products
    GROUP BY category
    ORDER BY total_stock DESC
""").pl().collect()

print(result)

import duckdb
import pandas as pd
import polars as pl
import numpy as np

# 初始化
con = duckdb.connect(":memory:")

# 1. 读取 CSV
con.sql("CREATE TABLE sales AS SELECT * FROM read_csv_auto('sales.csv')")

# 2. SQL 聚合
summary = con.sql("""
    SELECT 
        DATE_TRUNC('month', sale_date) AS month,
        category,
        SUM(amount) AS total_sales,
        COUNT(*) AS transactions
    FROM sales
    WHERE sale_date >= '2024-01-01'
    GROUP BY 1, 2
    ORDER BY 1, 2
""").df()

# 3. 注册 Polars 数据
products = pl.DataFrame({'id': [1,2,3], 'name': ['A','B','C']})
con.register('products', products)

# 4. 关联查询
enriched = con.sql("""
    SELECT s.*, p.name
    FROM summary s
    JOIN products p ON s.category = p.id
""").pl().collect()

print(f"处理了 {len(enriched)} 行数据")