pg_bestmatch

pig build pkg pg_bestmatch         # 构建 RPM / DEB 包

pig repo add pgsql -u          # 添加仓库并更新缓存

shared_preload_libraries = 'pg_bestmatch';

CREATE EXTENSION pg_bestmatch;

CREATE EXTENSION pg_bestmatch;
SET search_path TO public, bm_catalog;

wget https://huggingface.co/api/datasets/hazyresearch/LoCoV1-Documents/parquet/default/test/0.parquet -O documents.parquet
wget https://huggingface.co/api/datasets/hazyresearch/LoCoV1-Queries/parquet/default/test/0.parquet -O queries.parquet

import pandas as pd
from sqlalchemy import create_engine
import numpy as np
from psycopg2.extensions import register_adapter, AsIs

def adapter_numpy_float64(numpy_float64):
    return AsIs(numpy_float64)

def adapter_numpy_int64(numpy_int64):
    return AsIs(numpy_int64)

def adapter_numpy_float32(numpy_float32):
    return AsIs(numpy_float32)

def adapter_numpy_int32(numpy_int32):
    return AsIs(numpy_int32)

def adapter_numpy_array(numpy_array):
    return AsIs(tuple(numpy_array))

register_adapter(np.float64, adapter_numpy_float64)
register_adapter(np.int64, adapter_numpy_int64)
register_adapter(np.float32, adapter_numpy_float32)
register_adapter(np.int32, adapter_numpy_int32)
register_adapter(np.ndarray, adapter_numpy_array)

db_url = "postgresql://localhost:5432/pg_bestmatch_test"
engine = create_engine(db_url)

def load_documents():
    df = pd.read_parquet("documents.parquet")
    df.to_sql("documents", engine, if_exists='replace', index=False)

def load_queries():
    df = pd.read_parquet("queries.parquet")
    df['answer_pids'] = df['answer_pids'].apply(lambda x: str(x[0]))
    df.to_sql("queries", engine, if_exists='replace', index=False)

load_documents()
load_queries()

SELECT bm25_create('documents', 'passage', 'documents_passage_bm25', 0.75, 1.2);

ALTER TABLE documents ADD COLUMN embedding svector; -- pgvecto.rs 用户
ALTER TABLE documents ADD COLUMN embedding sparsevec; -- pgvector 用户

UPDATE documents SET embedding = bm25_document_to_svector('documents_passage_bm25', passage)::svector; -- pgvecto.rs 用户
UPDATE documents SET embedding = bm25_document_to_svector('documents_passage_bm25', passage, 'pgvector')::sparsevec; -- pgvector 用户

CREATE INDEX ON documents USING vectors (embedding svector_dot_ops); -- pgvecto.rs 用户
CREATE INDEX ON documents USING ivfflat (embedding sparsevec_ip_ops); -- pgvector 用户

ALTER TABLE queries ADD COLUMN embedding svector; -- pgvecto.rs 用户
ALTER TABLE queries ADD COLUMN embedding sparsevec; -- pgvector 用户

UPDATE queries SET embedding = bm25_query_to_svector('documents_passage_bm25', query)::svector; -- pgvecto.rs 用户
UPDATE queries SET embedding = bm25_query_to_svector('documents_passage_bm25', query, 'pgvector')::sparsevec; -- pgvector 用户

SELECT sum((array[answer_pids] = array(SELECT pid FROM documents WHERE queries.dataset = documents.dataset ORDER BY queries.embedding <#> documents.embedding LIMIT 1))::int) FROM queries;