Python API [Early Proposal]

[daxpryce]

Hey All!

I've been throwing together some ideas for a more pythonic API. These are still very early days and I'm soliciting feedback to try to get a sense for what mode of use people would actually like to see.

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT license.

import numpy as np

from abc import abstractmethod, ABC
from dataclasses import dataclass
from typing import overload, Any, Dict, List, Literal, Union
from os import PathLike

"""
Describes and documents the python diskann library
"""

Vector = List[Union[np.single, np.byte, np.ubyte]]


class DiskIndexBuilder:

    @overload
    def build(
        self,
        output_path: PathLike,
        vectors: List[Vector],
        distance_function: Literal["l2", "mips"],
        max_degree: int,
        build_complexity: int,
        search_memory_budget: int,
        build_memory_budget: int,
        index_prefix: str = "diskann",
        build_threads: int = 4,
        vector_target_product_quantization_bytes: int = 0,
        append_reorder_data: bool = False,
        build_product_quantization_bytes: int = 0,
        use_optimized_product_quantization: bool = False
    ):
        ...

    @overload
    def build(
        self,
        output_path: PathLike,
        vectors: np.ndarray,  # dtypes limited to np.single, np.byte, np.ubyte
        distance_function: Literal["l2", "mips"],
        max_degree: int,
        build_complexity: int,
        search_memory_budget: int,
        build_memory_budget: int,
        index_prefix: str = "diskann",
        build_threads: int = 4,
        vector_target_product_quantization_bytes: int = 0,
        append_reorder_data: bool = False,
        build_product_quantization_bytes: int = 0,
        use_optimized_product_quantization: bool = False
    ):
        ...

    @overload
    def build(
        self,
        output_path: PathLike,
        bin_path: PathLike,
        vector_dtype: Literal["float", "uint8", "int8"],
        distance_function: Literal["l2", "mips"],
        max_degree: int,
        build_complexity: int,
        search_memory_budget: int,
        build_memory_budget: int,
        index_prefix: str = "diskann",
        build_threads: int = 4,
        vector_target_product_quantization_bytes: int = 0,
        append_reorder_data: bool = False,
        build_product_quantization_bytes: int = 0,
        use_optimized_product_quantization: bool = False
    ):
        ...


class InMemoryIndexBuilder:
    @overload
    def build(
        self,
        output_path: PathLike,
        vectors: List[Vector],
        distance_function: Literal["l2", "mips"],
        max_degree: int,
        build_complexity: int,
        alpha: float,
        index_prefix: str = "diskann",
        build_threads: int = 4,
        build_product_quantization_bytes: int = 0,
        use_optimized_product_quantization: bool = False
    ):
        ...

    @overload
    def build(
        self,
        output_path: PathLike,
        vectors: np.ndarray,  # dtypes limited to np.single, np.byte, np.ubyte
        distance_function: Literal["l2", "mips"],
        max_degree: int,
        build_complexity: int,
        alpha: float,
        index_prefix: str = "diskann",
        build_threads: int = 4,
        build_product_quantization_bytes: int = 0,
        use_optimized_product_quantization: bool = False
    ):
        ...

    @overload
    def build(
        self,
        output_path: PathLike,
        bin_path: PathLike,
        vector_dtype: Literal["float", "uint8", "int8"],
        distance_function: Literal["l2", "mips"],
        max_degree: int,
        build_complexity: int,
        alpha: float,
        index_prefix: str = "diskann",
        build_threads: int = 4,
        build_product_quantization_bytes: int = 0,
        use_optimized_product_quantization: bool = False
    ):
        ...


@dataclass
class Neighbor:
    vector_id: int
    distance: float


class Index(ABC):

    @abstractmethod
    @overload
    def nearest(self, vector: Vector, k: int) -> List[Neighbor]:
        ...

    @overload
    def nearest(self, vectors: List[Vector], k: int) -> Dict[int, List[Neighbor]]:
        ...

    @overload
    def nearest(self, vectors: np.ndarray, k: int) -> Dict[int, List[Neighbor]]:
        ...

    @overload
    def nearest(self, vectors: Dict[Any, Vector], k: int) -> Dict[Any, List[Neighbor]]:
        ...


class DiskIndex(Index):
    def load(self): # TODO
        ...


class InMemoryIndex(Index):
    def load(self): # TODO
        ...


class MutableIndex(Index):
    def load(self): # TODO
        ...

    def delete(self, id: int): # TODO
        ...

    def upsert(self, id: int, vector: Vector): # TODO: It's not even clear we want a single upsert-like function or if we want standalone `insert` and `update` semantics instead.
        ...

    def optimize(self): # TODO: This is the method that tells DiskANN to recompute a new properly optimized index. It's still super unclear what this looks like outside of "something like this needs to exist"
        ...

[daxpryce]

I would also like to note that I have pulled out a 4th IndexBuilder overload that would instead shove a row-based, vector-projector function down into DiskANN lib for processing each point. I am not confident I know exactly how that would work and need to set up a toy project to explore it a bit more before I propose it.