[Feast] Feast 엔티티 키 직렬화 핫패스 최적화

PR 링크: feast-dev/feast#5981 상태: Merged | 변경: +712 / -30

들어가며

Feast의 Online Store에서 feature를 읽고 쓸 때마다 entity key의 직렬화/역직렬화가 발생한다. 실무에서는 90% 이상이 단일 entity key(예: user_id)를 사용하는데, 기존 코드는 이 경우에도 불필요한 sorted() 호출과 zip() + list comprehension을 수행했다. 역직렬화에서는 struct.unpack_from()이 매번 바이트 배열을 슬라이싱해 새 객체를 만들었다. 이 PR은 단일 entity fast path와 memoryview 기반 zero-copy 역직렬화로 핫패스를 최적화한다.

핵심 코드 분석

직렬화: 단일 entity fast path

Before:

def serialize_entity_key(entity_key, entity_key_serialization_version):
    if not entity_key.join_keys:
        sorted_keys = []
        sorted_values = []
    else:
        pairs = sorted(zip(entity_key.join_keys, entity_key.entity_values))
        sorted_keys = [k for k, _ in pairs]
        sorted_values = [v for _, v in pairs]

    output: List[bytes] = []
    for k in sorted_keys:
        output.append(struct.pack("<I", ValueType.STRING))
        if entity_key_serialization_version > 2:
            output.append(struct.pack("<I", len(k)))
        output.append(k.encode("utf8"))

After:

def serialize_entity_key(entity_key, entity_key_serialization_version):
    if not entity_key.join_keys:
        sorted_keys = []
        sorted_values = []
    elif len(entity_key.join_keys) == 1:
        # Fast path: single entity, no sorting needed
        sorted_keys = [entity_key.join_keys[0]]
        sorted_values = [entity_key.entity_values[0]]
    else:
        # Multi-entity: use sorting
        pairs = sorted(zip(entity_key.join_keys, entity_key.entity_values))
        sorted_keys = [k for k, _ in pairs]
        sorted_values = [v for _, v in pairs]

    output: List[bytes] = []
    if sorted_keys:
        encoded_keys = [k.encode("utf8") for k in sorted_keys]
        for i, k_encoded in enumerate(encoded_keys):
            output.append(struct.pack("<I", ValueType.STRING))
            if entity_key_serialization_version > 2:
                output.append(struct.pack("<I", len(k_encoded)))
            output.append(k_encoded)

단일 entity일 때 sorted(zip(...)) 대신 직접 인덱싱한다. O(n log n) 정렬이 O(1)로 줄어든다. 또한 키를 미리 한 번에 encode해서 len() 계산 시 바이트 길이와 문자열 길이의 불일치 문제도 방지한다.

prefix 직렬화도 동일한 패턴 적용

Before:

def serialize_entity_key_prefix(entity_keys, entity_key_serialization_version):
    sorted_keys = sorted(entity_keys)
    # ...

After:

def serialize_entity_key_prefix(entity_keys, entity_key_serialization_version):
    if len(entity_keys) == 1:
        sorted_keys = [entity_keys[0]]
    else:
        sorted_keys = sorted(entity_keys)
    # ...

역직렬화: memoryview zero-copy 슬라이싱

Before:

def deserialize_entity_key(serialized_entity_key, entity_key_serialization_version):
    offset = 0
    keys = []
    values = []

    num_keys = struct.unpack_from("<I", serialized_entity_key, offset)[0]
    offset += 4

    for _ in range(num_keys):
        key_type = struct.unpack_from("<I", serialized_entity_key, offset)[0]
        offset += 4
        key_length = struct.unpack_from("<I", serialized_entity_key, offset)[0]
        offset += 4
        key = struct.unpack_from(f"<{key_length}s", serialized_entity_key, offset)[0]
        keys.append(key.decode("utf-8").rstrip("\x00"))
        offset += key_length

After:

def deserialize_entity_key(serialized_entity_key, entity_key_serialization_version):
    buffer = memoryview(serialized_entity_key)
    pos = 0
    keys = []
    values = []

    num_keys = struct.unpack("<I", buffer[pos : pos + 4])[0]
    pos += 4

    for _ in range(num_keys):
        key_type, key_length = struct.unpack("<2I", buffer[pos : pos + 8])
        pos += 8
        key = struct.unpack(f"<{key_length}s", buffer[pos : pos + key_length])[0]
        keys.append(key.decode("utf-8").rstrip("\x00"))
        pos += key_length

memoryview를 사용해 바이트 슬라이싱 시 새 바이트 객체를 생성하지 않는다(zero-copy). 또한 key_type과 key_length를 한 번의 struct.unpack("<2I", ...)로 읽어 syscall 횟수를 줄인다. 경계 검사(bounds checking)도 추가해 잘못된 데이터에 대한 안전성을 높였다.

왜 이게 좋은가

1. 단일 entity 최적화: 전체 트래픽의 대다수를 차지하는 단일 entity 케이스에서 정렬 오버헤드를 완전히 제거한다.
2. Zero-copy 역직렬화: memoryview 슬라이싱은 새 바이트 객체를 할당하지 않아 GC 압박을 줄이고 처리 속도를 높인다.
3. 배치 언팩: struct.unpack("<2I", ...) 같은 배치 언팩으로 함수 호출 횟수를 줄인다.
4. 안전성 향상: 역직렬화에 경계 검사를 추가해 손상된 데이터에 대한 명확한 에러 메시지를 제공한다.

정리

Feature store의 online serving에서 entity key 직렬화는 매 요청마다 수십~수백 번 호출되는 핫패스다. 이 PR은 "90%가 단일 entity"라는 실무 패턴에 맞춰 fast path를 추가하고, memoryview로 zero-copy 역직렬화를 구현해 불필요한 객체 생성을 최소화했다. 벤치마크 테스트도 함께 추가되어 회귀를 감지할 수 있다.

참고 자료

• Feast Online Store 문서 -- Feast 온라인 feature 조회 개요
• Python memoryview 문서 -- Python memoryview zero-copy 슬라이싱

알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.