pgVector 기반 VectorDB 구축 및 효율적인 리소스(메모리,스토리지) 사용

이전에 임베딩 결과를 저장하는 데이터베이스로 vectorDB에 대해서 간단하게 알아보았다.

2024.08.04-Embedding을 저장하는 VectorDB 그리고 벡터 유사도 검색 Indexing

Embedding을 저장하는 VectorDB 그리고 벡터 유사도 검색 Indexing

 

 

이번 글에서는 로컬환경에서 vectorDB를 설치해 개발환경을 구축하고 테스트하려고 한다.

특히 vectorDB 관련 오프소스 중에서 pgVector를 사용하려고 한다.

postgres에서 제공해주기 때문에 접근하기 가장 쉬웠고, docker를 통해서도 간단하게 설치할수 있다. 

 

그리고 vectorDB에 쿼리를 직접 실행해서 확인해보려고 한다. (난 아무리도 이론보다는 직접 눈으로 봐야 직성이 풀린다)

테이블과 인덱스를 생성하고, 쿼리를 해서 유사도가 높은 벡터를 찾아보려고 한다. 부끄럽게도 이번 글을 계기로 그냥 가볍게 넘어갔던 <-> 연산자와 <=>연산자가 다른 방식이였단것을 알게 되었다..

 

또한 인덱스의 크기나 인덱스 생성시 메모리 사용량 등 리소스 성능 테스트도 확인해볼 예정이다. 실제로 운영환경에서 메모리 이슈가 있었기 때문이다. 이를 계기로 인덱스 생성시 메모리와 스토리지까지 디테일하게 고려해야 한다는것도 배웠다.

 

실제로 운영환경에서는 AWS RDS에서 Postgres에서 제공해주는 pgVector를 사용하고 있는데, 지난달 5월 AWS의 RDS for PostgreSql에서 pgvector 0.7.0 지원이 가능해졌다.  그래서 pgvector에서 RAG를 사용할수 있고, halfvec라와 sparsevec 라는 새로운 벡터 타입도 추가되었다.  halfvec는 2바이트 부동소수점으로 저장하는 반정밀 벡터이기 때문에, 앞으로 스토리지 용량이 부족해질 것으로 예상되는(?) 우리 환경에서 눈길을 끄는 변화였다. 그래서 이번 글의 마지막엔  vector타입과 halfvec 타입을 비교해서 halfvec 타입으로 변경할 경우 리소스 효율이 얼마나 있을지까지 확인해보려고 한다.

 

 

 

목차 

 

1. PgVector 설치 

1.1 Dockerfile 생성

1.2 docker-compse 생성

1.3 init_pgvector.sql 생성

1.4 Connection 확인 

 

2. Querying 

2.1 Data insert

2.2 Data upsert

2.3 Indexing

2.4 Select (Cosine연산, L2연산, 내적연산) 

 

3. 인덱스 튜닝 

3.1 인덱스 메모리 조절 (maintenance_work_mem)

3.2 인덱스 크기 확인

3.3 인덱스 스토리지 조절 (halfvec로 타입변경)

 

---

 

1. PgVector 설치

1.1 Dockerfile 생성

아래는 postgre의 최신버전을 베이스 이미지로 하고, pgvector 0.7.0 버전을 install 하는 Dockerfile을 작성한다. 

환경변수로 postgr user, password, db를 직접 지정할수 있다. 

# Use the official Postgres image as a base image
FROM postgres:latest

# Set environment variables for Postgres
ENV POSTGRES_USER=myuser
ENV POSTGRES_PASSWORD=mypassword
ENV POSTGRES_DB=mydb

# Install the build dependencies
USER root
RUN apt-get update && apt-get install -y \
    build-essential \
    git \
    postgresql-server-dev-all \
    && rm -rf /var/lib/apt/lists/*

# Clone, build, and install the pgvector extension -> upgrade 0.7.0 ver
RUN cd /tmp \
    && git clone --branch v0.7.0 https://github.com/pgvector/pgvector.git \
    && cd pgvector \
    && make \
    && make install

 

1.2 docker-compose 생성

위에서 생성한 Dockerfile과 같은 경로에 아래와 같은 docker-compose.yml 파일을생성한다.

기본적으로 5432 포트를 사용하고 아래에서 미리 작성할 sql 파일을 마운트한다. 

services:
  postgres:
    build: .
    ports:
      - "5432:5432"
    volumes:
      - ./data:/var/lib/postgresql/data
      - ./init_pgvector.sql:/docker-entrypoint-initdb.d/init_pgvector.sql
    environment:
      POSTGRES_USER: myuser
      POSTGRES_PASSWORD: mypassword
      POSTGRES_DB: mydb

 

1.3 init_pgvector.sql 생성

마찬가지로 또 같은 경로에 초기 쿼리를 작성해둘 init_pgvector.sql 파일을 생성한다. 

postgresql에서 extension으로 vector를 활성화하는 쿼리와 테스트할 테이블과 인덱스를 미리 쿼리로 작성해 둘수 있다.

(참고로 아래 예시에서 인덱스는 HNSW방식으로 생성하고, cosine 유사도 방식으로 distance를 계산한다)

-- Install the extension we just compiled

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE IF NOT EXISTS test_result
(
    keyword     varchar(30) not null constraint test_result_tmp_pkey1 primary key,
    embedding   vector(256) not null,
    modified_at timestamp   not null
);
CREATE INDEX IF NOT EXISTS idx__test_result__embedding
    ON test_result USING hnsw (embedding public.vector_cosine_ops);

 

1.4 Docker container 실행 

docker compose up 명령어를 통해 위에 설정한 환경에 맞춰 docker container를 실행시킨다. 

 

출처 : https://dev.to/stephenc222/how-to-use-postgresql-to-store-and-query-vector-embeddings-h4b 

 

1.5 Connection 확인 

마지막으로 실행된 pgvector에 아래와 같은 스크립트로 접속해볼수 있다. user,password, db는 설치할때 설정한 값이고, execute를 사용해서 쿼리 결과를 찍어볼수 있다. 참고로 아래 쿼리는 pgvectorDB의 extension 버전을 확인하는 쿼리로 print 결과는 0.7.0 이 나온다 

import psycopg2

# Establish a connection to the PostgreSQL database
conn = psycopg2.connect(
    user="myuser",
    password="mypassword",
    host="localhost",
    port=5432,
    database="mydb"
)
cur = conn.cursor()

# pgvector extension 버전 확인 쿼리
cur.execute("SELECT extversion FROM pg_extension WHERE extname = 'vector';")
print(cur.fetchone())

cur.close()
conn.close()

 

아 대신 python에서 postgresql client 패키지인 psycopg2를 사용하려면 요 패키지 install은 사전에 필요하다! 

pip3 install psycopg2-binary==2.9.9

 

또는 직접 database client(ex:DataGrip)에서 접속해서 설치했던 config(user,password,db)로 접속 해 볼수도 있다.

 

2. Querying 

2.1 Insert

벡터 쿼리를 하기 전에 임의의 데이터를 미리 테이블에 넣어둔다. 아래는 테스트하기 위한 dummy 데이터를 insert하는 스크립트이다. 

예를 들어, 현재 100개의 배치쿼리를 한번에 수행하면서, 최종적으로 1000개의 데이터를 Insert 하려고 한다.

import psycopg2
import numpy as np
from datetime import datetime
import random

# PostgreSQL 연결 설정
conn = psycopg2.connect(
    user="myuser",
    password="mypassword",
    host="localhost",
    port=5432,  # The port you exposed in docker-compose.yml
    database="mydb"
)
cursor = conn.cursor()

# 벡터 데이터 생성 함수
def generate_random_data(keyword:str, dimension: int):
    embedding = np.random.uniform(-1.0, 1.0, dimension).astype(np.float32)
    modified_at = datetime.now()
    return (keyword, embedding.tolist(), modified_at) 

def insert_batch_data(batch_data):
    insert_query = """
    INSERT INTO test_result (keyword, embedding, modified_at)
    VALUES (%s, %s, %s)
    """    
    # executemany를 사용하여 배치로 쿼리 실행
    cursor.executemany(insert_query, batch_data)
    conn.commit()

# 총 10만개의 데이터 생성 및 배치로 삽입
batch_size = 100
total_records = 1000
dim = 256 

for i in range(0, total_records, batch_size):
    batch_data = [
        generate_random_data(f"keyword_{j}", dim)
        for j in range(i, i + batch_size)
    ]
    print(f"insert batch data : {len(batch_data)}")
    insert_batch_data(batch_data)

# 연결 종료
cursor.close()
conn.close()

 

2.2 Upsert

postgresql 기반이기 때문에 변경된 내용만 업데이트해서 insert 하는 Updsert 문도 수행가능하다. 예를 들면 현재 테이블은 keyword가 primary key이기 때문에, Keyword가 같은 경우에만 embedding 값을 업데이트하고 싶다면 아래와 같이 수행할수 있다.

def upsert_data(batch_data):
    insert_query = """
    INSERT INTO test_result (keyword, embedding, modified_at)
    VALUES (%s, %s, %s)
    ON CONFLICT (keyword) DO UPDATE SET embedding = EXCLUDED.embedding;
    """    
    # executemany를 사용하여 배치로 쿼리 실행
    cursor.execute(insert_query, batch_data)
    conn.commit()

예를 들어, 위 쿼리를 통해 keyword값이 'keyword_0'인 row에 대해서 embedding값과 modified_at 값에 upsert를 수행했다. 

 

 

2.3 Indexing

pgvector는 neareast neighbor search를 통해 검색하기 떄문에 ANN을 하기 위해서는 인덱스를 추가해야한다. 현재 HNSW와 IVFFlat 방식의 인덱스를 지원하는데, 현업에서는 속도와 정확도가 높아 큰 데이터셋에 사용하기 좋은 HNSW를 방식을 사용하고 있다. 

 

그리고 인덱스는 속도와 정확도 사이에 tradeoff가 있다. HNSW의 옵션인 ef_constuction 값이 높을수록 인덱스 생성속도가 떨어지지만, 정확도가 높아진다. HNSW의 옵션의 기본값은 m 값이 16, ef_constuction 값이 64 이다. 현재 테스트 데이터의 양이 많지 않아서 이번 글에서는 아래와 같이 인덱스를 생성했다. 참고로 HNSW 방식에 대해서는 이전글에서 자세히 확인할수 있다. 

CREATE INDEX IF NOT EXISTS idx__test_result__embedding
ON test_result 
USING hnsw (embedding public.vector_cosine_ops)
WITH (m=8, ef_construction=32)

 

인덱스 생성하는 방식은 아래와 같이 3가지 방법이 있는데, 주로 코사인 유사도 방법을 사용한다. 참고로 그냥 지나갔을지도 모르겠지만 위에 init_pgvector.sql 부분에서 미리 생성한 인덱스 또한 코사인 방식이다.

# Euclidean 거리 (L2)  기반으로 인덱스 생성 
CREATE INDEX ON {table_name} USING hnsw (embedding vector_l2_ops)

# 내적값 (inner product) 기반으로 인덱스 생성
CREATE INDEX ON {table_name} USING hnsw (embedding vector_ip_ops)

# 코사인 (cosine) 기반으로 인덱스 생성
CREATE INDEX ON {table_name} USING hnsw (embedding vector_cosine_ops)

 

 

2.4 Select

일반적인 RDBMS의 select는 쿼리와 정확히 일치하는 행을 찾는 반면 VectoDB는 쿼리와 가장 유사한 행을 찾는다.

즉, 쿼리간의 검색은 두 쿼리간에 얼마나 유사한지, 즉 두 벡터간의 거리가 얼마나 가까운지를 계산한다. 검색방식 또한 인덱싱 기법과 동일해야하므로 마찬가지로 아래처럼 3가지 방법으로 쿼리를 검색할 수 있다. 

 

• vector1 <-> vector2: Euclidean 거리 (L2) 
• vector1 <#> vector2: 내적값 (inner product)
• vector1 <=> vector2:  Cosine 유사도

예를 들어, 쿼리벡터가 주어졌을때 해당 벡터와 유사한 값을 <코사인 유사도> 방식으로 찾고자 한다면 아래와 같은 쿼리가 될 것이다. 

select * FROM public.test_result
order by embedding <=> '{query_vector}'
limit 5

 

 

그리고 쿼리벡터 기준으로 얼마나 유사한지 score값까지 얻고 싶다면 1-cosine 값을 사용해 직접 구할수도 있다.

select keyword,embeeding, 1-(embeeding <=>'{query_vector}') as score
order by embedding <=> '{query_vector}'
limit 5

 

 

일단 기본적으로 vectorDB 쿼리와 인덱싱에 대해서 직접 쿼리를 날려보면서 알아보았다.

이정도면 pgvector를 써보는데 문제없을까 싶은데 조금더 엔지니어의 관점에서 성능튜닝 부분까지 알아보려고 한다. 

왜냐면 실제로 운영환경에서 몇가지 리소스 이슈가 있었기 때문이다.....!! 

 

 

 

3. 인덱스 튜닝

3.1 인덱스 메모리 조절 

인덱스 생성하는 작업은 메모리를 많이 사용한다.  그러다보니 적은 메모리로 큰 인덱스를 생성한다면 생성시간이 오래 걸릴 수 있다.   

 

인덱스 생성시 사용하는 메모리는 maintenace_work_mem 설정에 의해 컨트롤 된다.  

기본 메모리양은 64MB이며, 인덱스 생성 전에 이 값을 늘리면 인덱스 생성 작업을 더 빨리 완료할수 있다.  

 

실제로 운영환경에서 요 값을 8G로 아주 화끈하게 증가시켰었다. 약 100만개의 데이터에 대해서 인덱스 생성하는 해당 세션에서만 늘렸다.

SET MAINTENANCE_WORK_MEM='8GB'

 

그런데 이거 왠걸? 오히려 아래와 같은 메모리 이슈가 발생했었다. 

psycopg2.errors.OutOfMemory: could not resize shared memory segment “/PostgreSQL.1650105278” to 8586838528 bytes: Cannot allocate memory

 

maintenace_work_mem 값을 증가하시키면서  전체적으로 메모리 사용량이 증가해서 메모리 부족 이슈가 발생한 것 같았다.  메모리를 많이 써서 빨리 인덱스를 생성하고자 했던 급한 마음이 벌인 처참한 결과였다. 결론적으로 해당값을 조금 줄이니 위 이슈는 해결되었따. 

 

3.2 인덱스 크기 확인 

생성한 인덱스를 저장하기 위해 스토리지를 잡아먹는다. 그래서 적당한 크기의 인덱스를 생성하는것 또한 튜닝 포인트라고 생각한다. 

 

pg_stat_user_indexes 테이블에 사용자가 생성한 인덱스 관련 정보가 있다. 그래서 테이블 기준으로 생성한 인덱스의 크기와 이름을 확인할 수 있다. 또는 인덱스 크기가 큰 순으로 정렬해서 어떤 인덱스가 스토리지를 가장 많이 차지하는  확인할수 있다. 

위에서 생성한 인덱스의 이름은 idx_test_result_embeeding 로 1344KB 정도 사이즈가 되는걸 확인할수 있다.

 

추가로 이렇게만 확인하면, 인덱스의 크기가 어느정도인지 가늠이 잘 되지 않는다. 그래서 아래의 쿼리를 통해 인덱스의 크기를 테이블의 크기와 비교해서 확인해볼수 있다. 현재 테스트하기 위한 인덱스의 크기는 테이블의 크기에 비해 조금 큰 편인 것 같다.

 

 

3.3 인덱스 스토리지 조절 (halfvec로 타입변경)

최근 pgvector 7.0.0에서는 halfvec 타입을 지원한다. 확인해보니 실제로 기존 vector 타입으로 생성한 인덱스보다 크기가 줄어든다. 

기존이랑 같은 데이터를 넣고,  embedding의 타입만 vector-> halfvecf로 변경했다. 

create table if not exists public.test_result_half
(
    keyword     varchar(30) not null
        constraint test_result_half_tmp_pkey1
            primary key,
    embedding   halfvec(256) not null,
    modified_at timestamp   not null
);

create index if not exists idx__test_result_half__embedding
    on public.test_result_half using hnsw (embedding public.halfvec_cosine_ops);

 

그리고 기존에 있던 테이블의 인덱스 크기가 1344KB였는데, halfvec 타입으로 생성했을때 인덱스 크기가 824KB로 줄어든걸 확인할수 있다. 요 반정밀 데이터 타입을 사용해서 스토리지도 효율적으로 사용할수 있을것 같다.  

 

 

이상으로 도커기반으로 PgVector를 설치하는 방법부터, 직접 가상의 데이터를 넣어 검색 쿼리도 확인해보고, 인덱스 생성시 메모리와 스토리지 까지 고려해서 간단한 튜닝할수 있는 방법까지 알아보았다. 개인적으로 조금더 튜닝관련해서 찾아보고 싶은 부분이 더 많지만 끝도 없을것 같아서 이번 글은 여기에서 마무리하려고 한다. 

그리고 다음 글에서는 데이터베이스를 모니터링하는 방법에 대해서 작성하려고 한다. 요런 리소스 관련 부분을 알아보니 실제로 운영 중에 메모리가 튀는 부분이 있는지, 혹은 비효율적으로 사용하고 있는 부분이 있을까 싶어서  더 궁금하기 때문이다.

 

역시 데이터베이스는 끝이 없다..!