Elasticsearch Indexing & Vector Search

from elasticsearch import Elasticsearch
import numpy as np

# ES 클라이언트 연결
es = Elasticsearch("http://localhost:9200")
index_name = "image_search_index"

def create_index():
    """이미지 검색을 위한 Elasticsearch 인덱스 매핑(스키마) 정의"""
    mapping = {
        "mappings": {
            "properties": {
                "product_id": {"type": "keyword"},
                
                # 1. LSH 방식: 해시 서명들을 단순 문자열 배열(Keyword)로 저장
                # 검색 시 Term 쿼리로 완전 일치 매칭 수행 (역색인 활용)
                "lsh_signatures": {"type": "keyword"},
                
                # 2. HNSW 방식 (최신 ES 권장): 딥러닝 512차원 실수 벡터를 그대로 저장
                # index: true로 설정 시 내부적으로 HNSW 그래프 생성
                "cnn_vector": {
                    "type": "dense_vector",
                    "dims": 512,
                    "index": True, 
                    "similarity": "cosine" # 거리 계산 기준 (코사인 유사도)
                }
            }
        }
    }
    
    if not es.indices.exists(index=index_name):
        es.indices.create(index=index_name, body=mapping)
        print(f"인덱스 '{index_name}' 생성 완료.")

def index_document(product_id, lsh_sigs, vector):
    """추출된 데이터를 ES에 색인 (Indexing)"""
    doc = {
        "product_id": product_id,
        "lsh_signatures": lsh_sigs,       # 예: ["hash_A_101", "hash_B_010"]
        "cnn_vector": vector.tolist()     # Numpy 배열을 리스트로 변환
    }
    es.index(index=index_name, id=product_id, body=doc)

def search_by_lsh(query_signatures):
    """방식 1: LSH 토큰을 이용한 역색인 고속 검색 (Boolean Query)"""
    # 쿼리의 해시 토큰 중 '하나라도' 일치하는 문서들을 스코어 순으로 가져옴 (OR 조건)
    query = {
        "query": {
            "terms": {
                "lsh_signatures": query_signatures,
                "boost": 1.0
            }
        },
        "size": 10 # 상위 10개 상품 반환
    }
    response = es.search(index=index_name, body=query)
    return [hit["_source"]["product_id"] for hit in response["hits"]["hits"]]

def search_by_knn(query_vector):
    """방식 2: HNSW 그래프를 이용한 네이티브 벡터 검색 (k-NN Search)"""
    # k: 최종 반환할 결과 수, num_candidates: 각 샤드에서 HNSW로 탐색할 후보군 수
    query = {
        "knn": {
            "field": "cnn_vector",
            "query_vector": query_vector.tolist(),
            "k": 10,
            "num_candidates": 100
        },
        "_source": ["product_id"] # 반환할 필드 지정 (네트워크 대역폭 최적화)
    }
    response = es.search(index=index_name, body=query)
    return [hit["_source"]["product_id"] for hit in response["hits"]["hits"]]

# ==========================================
# 실행 시뮬레이션
# ==========================================
# 1. 스키마 세팅
create_index()

# 2. 데이터 인덱싱 (서버에 상품이 등록될 때 실행)
sample_vector = np.random.rand(512).astype(np.float32)
# LSH 테이블이 2개라고 가정하여 추출된 서명
sample_lsh_sigs = ["t1_10110", "t2_11100"] 

index_document("item_001", sample_lsh_sigs, sample_vector)
print("상품 데이터 색인 완료.")

# 3. 데이터 검색 (사용자가 쿼리 이미지를 올렸을 때 실행)
# 보통 실무에서는 1차로 LSH를 통해 100개로 후보군을 좁히고, 
# 2차로 해당 후보군에 대해서만 원본 벡터 코사인 유사도(Rescoring)를 계산하여 정확도를 극대화합니다.
lsh_results = search_by_lsh(["t1_10110", "t2_00000"])
print(f"LSH 기반 검색 결과: {lsh_results}")

knn_results = search_by_knn(sample_vector)
print(f"k-NN (HNSW) 기반 검색 결과: {knn_results}")