1. 벡터 검색에서의 다양성 확보 필요성
LLM을 활용한 RAG 시스템에서 단순히 벡터 유사도(Similarity)만으로 상위 K개의 청크를 가져올 경우, 특정 문서 하나에서 대부분의 결과가 추출될 위험이 있습니다. 이는 정보의 편향성을 초래하므로, 검색 결과의 다양성을 위해 각 문서(Document)당 최대 청크 노출 개수를 제한(예: 문서당 최대 2개)하는 로직이 필요합니다. 저자는 이를 위해 PostgreSQL 레벨에서 효율적으로 필터링하는 방법을 모색했습니다.
2. SQL 성능 최적화의 난관과 인덱스 활용
초기에 LLM(ChatGPT, Claude)을 통해 생성한 SQL 쿼리는 기능적으로는 작동했으나 성능 면에서 심각한 결함이 있었습니다.
- HNSW 인덱스 미사용: 쿼리가 복잡해지면서 PostgreSQL이 벡터 검색 전용 인덱스인 HNSW를 타지 않고 전체 테이블 스캔을 수행하는 문제가 발생했습니다.
- 불필요한 전체 정렬: 인덱스를 활용하지 못하고 전체 데이터를 정렬한 후 필터링을 시도하여 응답 시간이 수 초 이상으로 늘어났습니다.
저자는 EXPLAIN ANALYZE를 통해 실행 계획을 분석하고, 인덱스를 강제로 활용할 수 있는 쿼리 구조를 재설계했습니다.
3. CTE와 Window Function을 활용한 해결책
최종적으로 선택된 방법은 CTE(Common Table Expressions)와 윈도우 함수(Window Function)를 결합하는 것입니다.
- 내부 CTE (base): 우선 neighbor 젬을 사용하여 인덱스를 타는 최적화된 벡터 검색을 수행하고, inner_limit을 통해 필요한 것보다 넉넉한 양의 후보군을 먼저 추출합니다.
- 중간 CTE (partitioned_ranked): 추출된 후보군 내에서 ROW_NUMBER() OVER (PARTITION BY document_id ORDER BY neighbor_distance)를 사용하여 문서별 순위를 매깁니다.
- 최종 쿼리: 각 문서 내 순위가 설정된 max_per_document 이하인 결과만 필터링하여 반환합니다.
4. ActiveRecord 기반 구현 세부사항
저자는 이를 wrap_relation_for_max_per_interview라는 메서드로 모듈화했습니다. 이 메서드는 다음과 같은 특징을 가집니다.
- Arel 활용: 복잡한 윈도우 함수 구문을 Arel.sql을 통해 안전하게 삽입합니다.
- 메서드 체이닝: 기존 ActiveRecord::Relation 객체를 입력받아 CTE를 추가한 새로운 Relation을 반환하므로, Rails의 선언적 쿼리 작성 방식을 유지할 수 있습니다.
- 유연성: inner_limit 파라미터를 통해 오버페칭(Over-fetching) 정도를 조절하여, 필터링 후에도 충분한 결과값이 남도록 보장합니다.
5. 아키텍처적 고려사항: SQL vs Ruby
데이터베이스 레벨의 처리가 성능상 유리할 수 있지만, 저자는 Ruby 레이어에서의 처리 가능성도 열어두었습니다. 만약 검색 결과를 가져온 뒤 별도의 Re-ranker 모델을 적용해야 한다면, SQL에서 엄격하게 제한하는 것보다 Ruby에서 더 많은 데이터를 가져와 유연하게 처리하는 것이 나을 수 있기 때문입니다. 따라서 시스템의 요구사항에 따라 적절한 위치를 선택하는 것이 중요합니다.