[ PostgreSQL] PostgreSQL数组操作：为特征工程节省80%存储

-- 传统行存储方案（反面教材）
CREATE TABLE user_features_row (
    user_id BIGINT PRIMARY KEY,
    feature_1 FLOAT,
    feature_2 FLOAT,
    feature_3 FLOAT,
    -- ... 省略到 feature_100
    feature_100 FLOAT,
    updated_at TIMESTAMP DEFAULT NOW()
);

-- 数组化改造后的表结构
CREATE TABLE user_features_array (
    user_id BIGINT PRIMARY KEY,
    feature_vector FLOAT[],  -- 变长数组，存储100维特征
    feature_names TEXT[],    -- 可选：存储特征名映射
    updated_at TIMESTAMP DEFAULT NOW()
);

# 检查版本
psql -c "SELECT version();"

# 安装必要的扩展（超级用户执行）
psql -d your_database << 'EOF'
-- 数组操作增强扩展
CREATE EXTENSION IF NOT EXISTS intarray;      -- 整数数组操作
CREATE EXTENSION IF NOT EXISTS pg_stat_statements; -- 查询分析

-- 查看数组相关操作符
SELECT oprname, oprleft::regtype, oprright::regtype 
FROM pg_operator 
WHERE oprname IN ('@>', '<@', '&&', '||') 
AND oprleft::text LIKE '%[]';
EOF

-- 主特征表：存储稠密向量
CREATE TABLE ml.feature_vectors (
    entity_id BIGINT NOT NULL,           -- 用户/商品ID
    entity_type SMALLINT NOT NULL,       -- 1=用户, 2=商品, 3=店铺
    feature_version VARCHAR(32) NOT NULL, -- 特征版本号，如"v3.2.1"
    vector FLOAT[] NOT NULL,             -- 特征向量
    dim_count SMALLINT GENERATED ALWAYS AS (array_length(vector, 1)) STORED,
    created_at TIMESTAMP DEFAULT NOW(),
    
    PRIMARY KEY (entity_id, entity_type, feature_version)
) PARTITION BY LIST (entity_type);

-- 创建分区表（水平分片）
CREATE TABLE ml.feature_vectors_user PARTITION OF ml.feature_vectors
    FOR VALUES IN (1);
CREATE TABLE ml.feature_vectors_item PARTITION OF ml.feature_vectors
    FOR VALUES IN (2);

-- 特征元数据表：记录每个版本的特征含义
CREATE TABLE ml.feature_metadata (
    version VARCHAR(32) PRIMARY KEY,
    feature_names TEXT[] NOT NULL,       -- 特征名数组，与vector索引对应
    feature_dtypes CHAR(1)[] NOT NULL,   -- 'f'=float, 'i'=int, 'b'=bool
    description TEXT,
    created_by VARCHAR(64),
    created_at TIMESTAMP DEFAULT NOW()
);

-- 稀疏特征专用表（CSR格式存储）
CREATE TABLE ml.sparse_features (
    entity_id BIGINT NOT NULL,
    entity_type SMALLINT NOT NULL,
    version VARCHAR(32) NOT NULL,
    indices INT[],      -- 非零特征索引
    values FLOAT[],     -- 对应值
    PRIMARY KEY (entity_id, entity_type, version)
);

user_id,feat_0,feat_1,feat_2,...,feat_99
10001,0.23,0.0,1.56,...,0.89
10002,0.0,0.0,0.0,...,0.12

-- 步骤1：创建临时 staging 表
CREATE TEMP TABLE staging_features (
    user_id BIGINT,
    raw_features TEXT  -- 临时存储为文本，后续转换
) ON COMMIT DROP;

-- 步骤2：使用COPY快速导入（比INSERT快10倍）
COPY staging_features(user_id, raw_features)
FROM '/data/user_features.csv'
WITH (FORMAT csv, HEADER true, DELIMITER ',');

-- 步骤3：转换为数组并插入主表（批量操作）
INSERT INTO ml.feature_vectors (entity_id, entity_type, feature_version, vector)
SELECT 
    user_id,
    1,  -- 用户类型
    'v1.0.0',
    string_to_array(raw_features, ',')::FLOAT[]  -- 核心转换
FROM staging_features
WHERE raw_features IS NOT NULL 
  AND raw_features != '';

-- 创建GIN索引加速包含查询（可选）
CREATE INDEX idx_feature_vectors_gin ON ml.feature_vectors 
USING GIN (vector gin__float_ops);

# array_etl.py - 高性能特征导入
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
import psycopg2
from io import StringIO

class FeatureArrayLoader:
    def __init__(self, db_url: str):
        self.engine = create_engine(db_url)
        
    def numpy_to_postgres_array(self, arr: np.ndarray) -> str:
        """将numpy数组转换为PostgreSQL数组文本格式"""
        if arr.ndim == 1:
            return '{' + ','.join(map(str, arr)) + '}'
        return '{' + ','.join(self.numpy_to_postgres_array(row) for row in arr) + '}'
    
    def bulk_insert_features(self, df: pd.DataFrame, version: str, batch_size=10000):
        """
        批量插入特征数据，使用COPY协议
        df需要包含: entity_id, features(列名为feat_0, feat_1...)
        """
        # 将多列特征合并为数组
        feature_cols = [c for c in df.columns if c.startswith('feat_')]
        
        # 构造COPY数据流
        buffer = StringIO()
        for _, row in df.iterrows():
            features = row[feature_cols].values.astype(float)
            # 处理NaN（PostgreSQL数组不支持NaN，转为NULL或0）
            features = np.nan_to_num(features, nan=0.0)
            array_str = self.numpy_to_postgres_array(features)
            buffer.write(f"{row['entity_id']}\t1\t{version}\t{array_str}\n")
        
        buffer.seek(0)
        
        # 使用COPY FROM
        conn = psycopg2.connect(self.engine.url)
        cursor = conn.cursor()
        cursor.copy_from(
            buffer, 
            'ml.feature_vectors',
            columns=('entity_id', 'entity_type', 'feature_version', 'vector'),
            sep='\t'
        )
        conn.commit()
        cursor.close()
        conn.close()
        
        print(f"Inserted {len(df)} rows into feature_vectors")

# 使用示例
if __name__ == "__main__":
    loader = FeatureArrayLoader("postgresql://user:pass@localhost/db")
    
    # 读取Spark生成的Parquet（假设已转为pandas）
    df = pd.read_parquet("user_features.parquet")
    
    # 分批导入避免内存爆炸
    for i in range(0, len(df), 10000):
        batch = df.iloc[i:i+10000]
        loader.bulk_insert_features(batch, version="v2.1.0")

-- 1. 提取单维度特征（比行存储的列访问慢，但省空间）
SELECT 
    entity_id,
    vector[1] as first_feature,      -- 数组下标从1开始！
    vector[10:20] as feature_slice,  -- 切片操作
    array_length(vector, 1) as dim   -- 获取维度
FROM ml.feature_vectors
WHERE entity_id = 10001;

-- 2. 数组聚合：计算全局统计
SELECT 
    feature_version,
    AVG(vector[1]) as f1_mean,
    STDDEV(vector[1]) as f1_std,
    PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY vector[5]) as f5_median
FROM ml.feature_vectors
GROUP BY feature_version;

-- 3. 数组解构：展开为行（用于与传统ETL兼容）
SELECT entity_id, feature_version, 
       unnest(vector) as feature_value,
       generate_subscripts(vector, 1) as feature_idx
FROM ml.feature_vectors
WHERE entity_id = 10001;

-- 1. 向量相似度计算（余弦相似度）
CREATE OR REPLACE FUNCTION cosine_similarity(a FLOAT[], b FLOAT[])
RETURNS FLOAT AS $$
DECLARE
    dot FLOAT := 0;
    norm_a FLOAT := 0;
    norm_b FLOAT := 0;
    i INT;
BEGIN
    IF array_length(a, 1) != array_length(b, 1) THEN
        RAISE EXCEPTION 'Array dimensions do not match';
    END IF;
    
    FOR i IN 1..array_length(a, 1) LOOP
        dot := dot + a[i] * b[i];
        norm_a := norm_a + a[i]^2;
        norm_b := norm_b + b[i]^2;
    END LOOP;
    
    IF norm_a = 0 OR norm_b = 0 THEN
        RETURN 0;
    END IF;
    
    RETURN dot / (sqrt(norm_a) * sqrt(norm_b));
END;
$$ LANGUAGE plpgsql IMMUTABLE;

-- 使用：找到与目标用户最相似的10个用户
WITH target AS (
    SELECT vector as target_vec
    FROM ml.feature_vectors
    WHERE entity_id = 10001 AND feature_version = 'v1.0.0'
)
SELECT 
    f.entity_id,
    cosine_similarity(f.vector, t.target_vec) as sim_score
FROM ml.feature_vectors f, target t
WHERE f.entity_id != 10001
ORDER BY sim_score DESC
LIMIT 10;

-- 2. 批量向量运算（使用数组操作符）
SELECT 
    entity_id,
    vector + ARRAY[0.1, -0.1, 0.0, ...]::FLOAT[] as shifted_vector,  -- 平移
    vector * 2.0 as scaled_vector,                                    -- 缩放
    sqrt((vector ^ 2)) as l2_norm                                     -- L2范数（自定义操作符）
FROM ml.feature_vectors
WHERE feature_version = 'v1.0.0';

-- 启用JIT编译（PostgreSQL 11+）
SET jit = on;
SET jit_above_cost = 10000;  -- 成本超过此阈值启用JIT

-- 分析查询计划
EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON)
SELECT cosine_similarity(vector, (SELECT vector FROM ml.feature_vectors WHERE entity_id=1))
FROM ml.feature_vectors;

-- 创建分区表（按时间分区，自动归档）
CREATE TABLE user_features (
    user_id BIGINT,
    feature_date DATE,  -- 分区键
    version VARCHAR(16),
    vector FLOAT[],
    PRIMARY KEY (user_id, feature_date, version)
) PARTITION BY RANGE (feature_date);

-- 自动创建未来3个月分区
DO $$
DECLARE
    start_date DATE := CURRENT_DATE;
    end_date DATE := CURRENT_DATE + INTERVAL '3 months';
    cur_date DATE := start_date;
BEGIN
    WHILE cur_date < end_date LOOP
        EXECUTE format(
            'CREATE TABLE IF NOT EXISTS user_features_%s 
             PARTITION OF user_features
             FOR VALUES FROM (%L) TO (%L)',
            TO_CHAR(cur_date, 'YYYYMM'),
            cur_date,
            cur_date + INTERVAL '1 month'
        );
        cur_date := cur_date + INTERVAL '1 month';
    END LOOP;
END $$;

-- 旧分区自动压缩（PostgreSQL 14+）
ALTER TABLE user_features_202401 SET (compression = 'zstd');

# feature_store.py - 生产级特征服务
import asyncpg
import numpy as np
from typing import List, Dict, Optional
import hashlib
import asyncio

class PGFeatureStore:
    def __init__(self, dsn: str, pool_size: int = 20):
        self.dsn = dsn
        self.pool = None
        self._cache = {}  # 本地LRU缓存
        
    async def initialize(self):
        self.pool = await asyncpg.create_pool(
            self.dsn,
            min_size=5,
            max_size=20,
            command_timeout=60
        )
        
    async def get_features_batch(
        self, 
        entity_ids: List[int], 
        version: str,
        as_numpy: bool = True
    ) -> Dict[int, np.ndarray]:
        """
        批量获取特征，使用unnest实现高效IN查询
        """
        async with self.pool.acquire() as conn:
            # 使用unnest将Python列表转为SQL数组
            rows = await conn.fetch(
                """
                SELECT entity_id, vector 
                FROM ml.feature_vectors
                WHERE entity_id = ANY($1::bigint[])
                  AND feature_version = $2
                  AND entity_type = 1
                """,
                entity_ids,
                version
            )
            
        result = {}
        for row in rows:
            vec = np.array(row['vector'], dtype=np.float32) if as_numpy else row['vector']
            result[row['entity_id']] = vec
            
        # 记录缺失
        missing = set(entity_ids) - set(result.keys())
        if missing:
            print(f"Missing features for entities: {missing}")
            
        return result
    
    async def compute_similar_users(
        self,
        target_id: int,
        top_k: int = 100,
        min_similarity: float = 0.8
    ) -> List[Dict]:
        """
        使用pgvector扩展（可选）或纯SQL计算相似用户
        """
        async with self.pool.acquire() as conn:
            # 方法1：纯PostgreSQL数组计算（无额外依赖）
            rows = await conn.fetch(
                """
                WITH target AS (
                    SELECT vector as target_vec
                    FROM ml.feature_vectors
                    WHERE entity_id = $1 AND feature_version = 'v1.0.0'
                )
                SELECT 
                    f.entity_id,
                    ($2 <-> f.vector) as distance,  -- L2距离，使用<->操作符
                    1 - (($2 <-> f.vector) / 2) as similarity  -- 归一化相似度
                FROM ml.feature_vectors f, target t
                WHERE f.entity_id != $1
                  AND f.feature_version = 'v1.0.0'
                ORDER BY distance
                LIMIT $3
                """,
                target_id,
                await self._get_vector(conn, target_id),
                top_k
            )
            
        return [dict(row) for row in rows]
    
    async def _get_vector(self, conn, entity_id: int) -> list:
        row = await conn.fetchrow(
            "SELECT vector FROM ml.feature_vectors WHERE entity_id = $1",
            entity_id
        )
        return row['vector'] if row else None

# 性能测试
async def benchmark():
    store = PGFeatureStore("postgresql://localhost/feature_db")
    await store.initialize()
    
    # 模拟批量推理请求
    batch_sizes = [1, 10, 100, 1000]
    
    for batch_size in batch_sizes:
        import time
        ids = list(range(1, batch_size + 1))
        
        start = time.time()
        for _ in range(100):
            await store.get_features_batch(ids, "v1.0.0")
        elapsed = (time.time() - start) / 100 * 1000  # ms
        
        print(f"Batch size {batch_size}: {elapsed:.2f}ms avg")

if __name__ == "__main__":
    asyncio.run(benchmark())

-- ❌ 错误1：混合类型
SELECT ARRAY[1, 2.5, 3];  -- 失败：无法混合int和float

-- ✅ 修复：显式类型转换
SELECT ARRAY[1::float, 2.5, 3::float];

-- ❌ 错误2：越界访问返回NULL而非报错
SELECT (ARRAY[1,2,3])[10];  -- 返回NULL，可能引发静默错误

-- ✅ 修复：使用array_get with边界检查
CREATE OR REPLACE FUNCTION safe_array_get(arr ANYARRAY, idx INT, default_val ANYELEMENT)
RETURNS ANYELEMENT AS $$
BEGIN
    IF idx < 1 OR idx > array_length(arr, 1) THEN
        RETURN default_val;
    END IF;
    RETURN arr[idx];
END;
$$ LANGUAGE plpgsql IMMUTABLE;

-- ❌ 错误3：大数组的WAL日志膨胀
-- 更新100万元素的数组会产生大量WAL

-- ✅ 修复：使用TOAST压缩+批量更新
ALTER TABLE ml.feature_vectors ALTER COLUMN vector SET STORAGE EXTERNAL;
-- 或使用逻辑复制槽过滤大字段

-- 监控数组表膨胀
SELECT 
    schemaname,
    tablename,
    pg_size_pretty(pg_total_relation_size(schemaname||'.'||tablename)) as total_size,
    pg_size_pretty(pg_relation_size(schemaname||'.'||tablename)) as table_size,
    pg_size_pretty(pg_indexes_size(schemaname||'.'||tablename)) as index_size
FROM pg_tables
WHERE tablename LIKE '%feature%'
ORDER BY pg_total_relation_size(schemaname||'.'||tablename) DESC;

-- 检查TOAST压缩率
SELECT 
    relname,
    pg_size_pretty(pg_relation_size(oid)) as main_size,
    pg_size_pretty(pg_relation_size(reltoastrelid)) as toast_size,
    CASE WHEN pg_relation_size(reltoastrelid) > 0 
         THEN round(100.0 * pg_relation_size(reltoastrelid) / pg_relation_size(oid), 2)
         ELSE 0 
    END as toast_ratio
FROM pg_class
WHERE relname = 'feature_vectors';

-- 自动维护任务
CREATE OR REPLACE FUNCTION maintain_feature_tables()
RETURNS void AS $$
BEGIN
    -- 更新统计信息
    ANALYZE ml.feature_vectors;
    
    -- 清理旧版本（保留最近3个）
    DELETE FROM ml.feature_vectors
    WHERE (entity_id, entity_type, created_at) NOT IN (
        SELECT entity_id, entity_type, MAX(created_at)
        FROM ml.feature_vectors
        GROUP BY entity_id, entity_type
        HAVING COUNT(*) > 3
    );
    
    -- 重建索引（如果膨胀严重）
    REINDEX INDEX CONCURRENTLY idx_feature_vectors_pk;
END;
$$ LANGUAGE plpgsql;