infinity-sdk 0.6.3__py3-none-any.whl

infinity/__init__.py

@@ -0,0 +1,34 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# import importlib.metadata
+#
+# __version__ = importlib.metadata.version("infinity_sdk")
+
+import logging
+import os
+# import pkg_resources
+# __version__ = pkg_resources.get_distribution("infinity_sdk").version
+
+from infinity.common import URI, NetworkAddress, LOCAL_HOST, LOCAL_INFINITY_PATH, InfinityException
+from infinity.infinity import InfinityConnection
+from infinity.remote_thrift.infinity import RemoteThriftInfinityConnection
+from infinity.errors import ErrorCode
+
+
+def connect(uri=LOCAL_HOST, logger: logging.Logger = None) -> InfinityConnection:
+    if isinstance(uri, NetworkAddress):
+        return RemoteThriftInfinityConnection(uri, logger)
+    else:
+        raise InfinityException(ErrorCode.INVALID_SERVER_ADDRESS, f"Unknown uri: {uri}")

infinity/common.py

@@ -0,0 +1,124 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pathlib import Path
+from typing import Union
+from dataclasses import dataclass, field
+import numpy as np
+
+
+class NetworkAddress:
+    def __init__(self, ip, port):
+        self.ip = ip
+        self.port = port
+
+    def __str__(self):
+        return f'IP: {self.ip}, Port: {self.port}'
+
+
+@dataclass
+class SparseVector:
+    indices: list[int]
+    values: Union[list[float], list[int], None] = None
+
+    def __post_init__(self):
+        assert (self.values is None) or (len(self.indices) == len(self.values))
+
+    def to_dict_old(self):
+        d = {"indices": self.indices}
+        if self.values is not None:
+            d["values"] = self.values
+        return d
+
+    def to_dict(self):
+        if self.values is None:
+            raise ValueError("SparseVector.values is None")
+        result = {}
+        for i, v in zip(self.indices, self.values):
+            result[str(i)] = v
+        return result
+
+    @staticmethod
+    def from_dict(d):
+        return SparseVector(d["indices"], d.get("values"))
+
+    def __str__(self):
+        return f"SparseVector(indices={self.indices}{'' if self.values is None else f', values={self.values}'})"
+
+    def __repr__(self):
+        return str(self)
+
+
+@dataclass
+class Array:
+    elements: list = field(default_factory=list)
+
+    def append(self, element):
+        self.elements.append(element)
+
+    def __init__(self, *args):
+        self.elements = list(args)
+
+    def __str__(self):
+        return f"Array({', '.join(str(e) for e in self.elements)})"
+
+    def __repr__(self):
+        return str(self)
+
+
+@dataclass
+class FDE:
+    """Fixed Dimensional Encoding function call for insert operations."""
+    tensor_data: list  # 2D array of numbers
+    target_dimension: int  # Target embedding dimension
+
+    def __init__(self, tensor_data: list, target_dimension: int):
+        self.tensor_data = tensor_data
+        self.target_dimension = target_dimension
+
+    def __str__(self):
+        return f"FDE(tensor_data={self.tensor_data}, target_dimension={self.target_dimension})"
+
+    def __repr__(self):
+        return str(self)
+
+
+URI = Union[NetworkAddress, Path]
+VEC = Union[list, np.ndarray, FDE]
+INSERT_DATA = dict[str, Union[str, int, float, list[Union[int, float]]], SparseVector, dict, Array, FDE]
+
+LOCAL_HOST = NetworkAddress("127.0.0.1", 23817)
+
+# test embedded_infinity
+LOCAL_INFINITY_PATH = "/var/infinity"
+
+
+class ConflictType(object):
+    Ignore = 0
+    Error = 1
+    Replace = 2
+
+
+class SortType(object):
+    Asc = 0
+    Desc = 1
+
+
+class InfinityException(Exception):
+    def __init__(self, error_code=0, error_msg=None):
+        self.error_code = error_code
+        self.error_msg = error_msg
+
+
+DEFAULT_MATCH_VECTOR_TOPN = 10
+DEFAULT_MATCH_SPARSE_TOPN = 10

infinity/connection_pool.py

@@ -0,0 +1,59 @@
+# Copyright(C) 2024 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from threading import Lock
+import infinity
+from infinity.common import NetworkAddress
+import logging
+
+
+class ConnectionPool(object):
+    def __init__(self, uri=NetworkAddress("127.0.0.1", 23817), max_size=16):
+        self.uri_ = uri
+        self.max_size_ = max_size
+        self.free_pool_ = []
+        self.lock_ = Lock()
+        for i in range(max_size):
+            self._create_conn()
+
+    def _del__(self):
+        self.destroy()
+
+    def _create_conn(self):
+        infinity_coon = infinity.connect(self.uri_)
+        self.free_pool_.append(infinity_coon)
+
+    def get_conn(self):
+        with self.lock_:
+            if (len(self.free_pool_) == 0):
+                self._create_conn()
+            conn = self.free_pool_.pop()
+            logging.debug("get_conn")
+            return conn
+
+    def release_conn(self, conn):
+        """
+        Note: User is allowed to release a connection not created by ConnectionPool, or not releasing(due to exception or some other reasons) a connection created by ConnectionPool.
+        """
+        with self.lock_:
+            if (self.free_pool_.count(conn)):
+                raise Exception("the connection has been released")
+            if (len(self.free_pool_) < self.max_size_):
+                self.free_pool_.append(conn)
+            logging.debug("release_conn")
+
+    def destroy(self):
+        for conn in iter(self.free_pool_):
+            conn.disconnect()
+        self.free_pool_.clear()

infinity/db.py

@@ -0,0 +1,48 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+
+
+class Database(ABC):
+
+    @abstractmethod
+    def create_table(self, table_name, schema, options):
+        pass  # implement create table logic here
+
+    @abstractmethod
+    def drop_table(self, table_name):
+        pass  # implement drop table logic here
+
+    @abstractmethod
+    def list_tables(self):
+        pass  # implement list tables logic here
+
+    @abstractmethod
+    def show_table(self, table_name):
+        pass  # implement describe table logic here
+
+    @abstractmethod
+    def get_table(self, table_name):
+        pass  # implement get table logic here
+
+    @abstractmethod
+    def create_table_snapshot(self, snapshot_name: str, table_name: str):
+        pass
+
+    @abstractmethod
+    def restore_table_snapshot(self, snapshot_name: str):
+        pass
+    
+    

infinity/errors.py

@@ -0,0 +1,183 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import IntEnum
+
+
+class ErrorCode(IntEnum):
+    OK = 0,
+
+    INVALID_TIME_INFO = 1001,
+    EMPTY_CONFIG_PARAMETER = 1002,
+    MISMATCH_VERSION = 1003,
+    INVALID_TIMEZONE = 1004,
+    INVALID_BYTE_SIZE = 1005,
+    INVALID_IP_ADDR = 1006,
+    INVALID_LOG_LEVEL = 1007,
+    INVALID_CONFIG = 1008
+
+    WRONG_PASSWD = 2001,
+    INSUFFICIENT_PRIVILEGE = 2002,
+    UNSUPPORTED_VERSION_INDEX = 2003,
+    CLIENT_VERSION_MISMATCH = 2004,
+    INFINITY_IS_STARTING = 2007,
+    INFINITY_IS_INITING = 2008,
+
+    INVALID_USERNAME = 3001,
+    INVALID_PASSWD = 3002,
+    INVALID_IDENTIFIER_NAME = 3003,
+    INVALID_TABLE_NAME = 3004,
+    INVALID_COLUMN_NAME = 3005,
+    INVALID_INDEX_NAME = 3006,
+    INVALID_COLUMN_DEFINITION = 3007,
+    INVALID_TABLE_DEFINITION = 3008,
+    INVALID_INDEX_DEFINITION = 3009,
+    DATA_TYPE_MISMATCH = 3010,
+    NAME_TOO_LONG = 3011,
+    RESERVED_NAME = 3012,
+    SYNTAX_ERROR = 3013,
+    INVALID_PARAMETER_VALUE = 3014,
+    DUPLICATE_USER = 3015,
+    DUPLICATE_DATABASE_NAME = 3016,
+    DUPLICATE_TABLE_NAME = 3017,
+    DUPLICATE_INDEX_NAME = 3018,
+    DUPLICATE_INDEX_ON_COLUMN = 3019,
+    NO_SUCH_USER = 3020,
+    DB_NOT_EXIST = 3021,
+    TABLE_NOT_EXIST = 3022,
+    INDEX_NOT_EXIST = 3023,
+    COLUMN_NOT_EXIST = 3024,
+    AGG_NOT_ALLOW_IN_WHERE_CLAUSE = 3025,
+    COLUMN_IN_SELECT_NOT_IN_GROUP_BY = 3026,
+    NO_SUCH_SYSTEM_VAR = 3027,
+    INVALID_SYSTEM_VAR_VALUE = 3028,
+    SYSTEM_VAR_READ_ONLY = 3029,
+    FUNCTION_NOT_FOUND = 3030,
+    SPECIAL_FUNCTION_NOT_FOUND = 3031,
+    NOT_SUPPORTED = 3032,
+    DROPPING_USING_DB = 3033,
+    SESSION_NOT_FOUND = 3034,
+    RECURSIVE_AGG = 3035,
+    FUNCTION_ARGS_ERROR = 3036,
+    IMPORT_FILE_FORMAT_ERROR = 3037,
+    DATA_NOT_EXIST = 3038,
+    COLUMN_COUNT_MISMATCH = 3039,
+    EMPTY_DB_NAME = 3040,
+    EMPTY_TABLE_NAME = 3041,
+    EMPTY_COLUMN_NAME = 3042,
+    EMPTY_INDEX_NAME = 3043,
+    EXCEED_DB_NAME_LENGTH = 3044,
+    EXCEED_TABLE_NAME_LENGTH = 3045,
+    EXCEED_COLUMN_NAME_LENGTH = 3046,
+    EXCEED_INDEX_NAME_LENGTH = 3047,
+    NO_COLUMN_DEFINED = 3048,
+    NOT_SUPPORTED_TYPE_CONVERSION = 3049,
+    EMPTY_SELECT_FIELDS = 3050,
+    INVALID_DATA_TYPE = 3051,
+    PARSE_MATCH_EXPR_FAILED = 3052,
+    FTS_INDEX_NOT_EXIST = 3053,
+    UNKNOWN_FTS_FAULT = 3054,
+    INVALID_CONSTRAINT_TYPE = 3055,
+    INVALID_KNN_DISTANCE_TYPE = 3056,
+    INVALID_EMBEDDING_DATA_TYPE = 3057,
+    INVALID_CONSTANT_TYPE = 3058,
+    INVALID_PARSED_EXPR_TYPE = 3059,
+    INVALID_INDEX_TYPE = 3060,
+    INVALID_INDEX_PARAM = 3061,
+    LACK_INDEX_PARAM = 3062,
+    INVALID_FILTER_EXPRESSION = 3063,
+    MUTIPLE_FUNCTION_MATCHED = 3064,
+    INSERT_WITHOUT_VALUES = 3065,
+    INVALID_CONFLICT_TYPE = 3066,
+    INVALID_JSON_FORMAT = 3067,
+    DUPLICATE_COLUMN_NAME = 3068,
+    INVALID_EXPRESSION = 3069,
+    SEGMENT_NOT_EXIST = 3070,
+    AGGREGATE_FUNCTION_WITH_EMPTY_ARGS = 3071,
+    BLOCK_NOT_EXIST = 3072,
+    INVALID_TOPK_TYPE = 3073,
+    INVALID_CREATE_OPTION = 3074,
+    INVALID_DROP_OPTION = 3075,
+    INVALID_COMMAND = 3076,
+    ANALYZER_NOT_FOUND = 3077,
+    NOT_SUPPORTED_ANALYZER = 3078,
+    INVALID_ANALYZER_NAME = 3079,
+    INVALID_ANALYZER_FILE = 3080,
+    INVALID_EXPLAIN_TYPE = 3081,
+    CHUNK_NOT_EXIST = 3082,
+    NAME_MISMATCHED = 3083,
+    TRANSACTION_NOT_FOUND = 3084,
+    INVALID_DATABASE_INDEX = 3085,
+    INVALID_TABLE_INDEX = 3086,
+    FUNCTION_IS_DISABLE = 3087,
+    NOT_FOUND = 3088,
+    ERROR_INIT = 3089,
+    FILE_IS_OPEN = 3090,
+    UNKNOWN = 3091,
+    INVALID_QUERY_OPTION = 3092,
+
+    TXN_ROLLBACK = 4001,
+    TXN_CONFLICT = 4002,
+
+    DISK_FULL = 5001,
+    OUT_OF_MEMORY = 5002,
+    TOO_MANY_CONNECTIONS = 5003,
+    CONFIGURATION_LIMIT_EXCEED = 5004,
+    QUERY_IS_TOO_COMPLEX = 5005,
+    FAIL_TO_GET_SYS_INFO = 5006,
+
+    QUERY_CANCELLED = 6001,
+    QUERY_NOT_SUPPORTED = 6002,
+    CLIENT_CLOSE = 6003,
+
+    DISK_IO_ERROR = 7001,
+    DUPLICATED_FILE = 7002,
+    CONFIG_FILE_ERROR = 7003,
+    LOCK_FILE_EXISTS = 7004,
+    CATALOG_CORRUPTED = 7005,
+    DATA_CORRUPTED = 7006,
+    INDEX_CORRUPTED = 7007,
+    FILE_NOT_FOUND = 7008,
+    DIR_NOT_FOUND = 7009,
+    DATA_IO_ERROR = 7010,
+    UNEXPECTED_ERROR = 7011,
+    PARSER_ERROR = 7012,
+    MMAP_FILE_ERROR = 7013,
+    MUNMAP_FILE_ERROR = 7014,
+    INVALID_FILE_FLAG = 7015,
+    INVALID_SERVER_ADDRESS = 7016,
+    FAIL_TO_FUN_PYTHON = 7017,
+    CANT_CONNECT_SERVER = 7018,
+    NOT_EXIST_NODE = 7019,
+    DUPLICATE_NODE = 7020,
+    CANT_CONNECT_LEADER = 7021,
+    MINIO_INVALID_ACCESS_KEY = 7022,
+    MINIO_BUCKET_NOT_EXISTS = 7023,
+    INVALID_STORAGE_TYPE = 7024,
+    NOT_REGISTERED = 7025,
+    CANT_SWITCH_ROLE = 7026,
+    TOO_MANY_FOLLOWER = 7027,
+    TOO_MANY_LEARNER = 7028,
+    CHECKPOINTING = 7029,
+
+    INVALID_ENTRY = 8001,
+    DUPLICATE_ENTRY = 8002
+    NOT_FOUND_ENTRY = 8003,
+    EMPTY_ENTRY_LIST = 8004,
+    NO_WAL_ENTRY_FOUND = 8005,
+    WRONG_CHECKPOINT_TYPE = 8006,
+    INVALID_NODE_ROLE = 8007,
+    INVALID_NODE_STATUS = 8008,
+    NODE_INFO_UPDATED = 8009,
+    NODE_NAME_MISMATCH = 8010

infinity/index.py

@@ -0,0 +1,103 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import Enum
+
+import infinity.remote_thrift.infinity_thrift_rpc.ttypes as ttypes
+from infinity.common import InfinityException
+from infinity.errors import ErrorCode
+
+
+class IndexType(Enum):
+    IVF = 1
+    Hnsw = 2
+    FullText = 3
+    Secondary = 4
+    EMVB = 5
+    BMP = 6
+    DiskAnn = 7
+
+    def to_ttype(self):
+        match self:
+            case IndexType.IVF:
+                return ttypes.IndexType.IVF
+            case IndexType.Hnsw:
+                return ttypes.IndexType.Hnsw
+            case IndexType.FullText:
+                return ttypes.IndexType.FullText
+            case IndexType.Secondary:
+                return ttypes.IndexType.Secondary
+            case IndexType.EMVB:
+                return ttypes.IndexType.EMVB
+            case IndexType.BMP:
+                return ttypes.IndexType.BMP
+            case IndexType.DiskAnn:
+                return ttypes.IndexType.DiskAnn
+            case _:
+                raise InfinityException(ErrorCode.INVALID_INDEX_TYPE, "Unknown index type")
+
+
+class InitParameter:
+    def __init__(self, param_name: str, param_value: str):
+        self.param_name = param_name
+        self.param_value = param_value
+
+    def __str__(self):
+        return f"InitParameter({self.param_name}, {self.param_value})"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def to_ttype(self):
+        return ttypes.InitParameter(self.param_name, self.param_value)
+
+
+class IndexInfo:
+    def __init__(self, column_name: str, index_type: IndexType, params: dict = None):
+        self.column_name = column_name
+        self.index_type = index_type
+        if params is not None:
+            if isinstance(params, dict):
+                self.params = params
+            else:
+                raise InfinityException(ErrorCode.INVALID_INDEX_PARAM, f"{params} should be dictionary type")
+        else:
+            self.params = None
+
+    def __str__(self):
+        return f"IndexInfo({self.column_name}, {self.index_type}, {self.params})"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __eq__(self, other):
+        return self.column_name == other.index_name and self.index_type == other.index_type and self.params == other.params
+
+    def __hash__(self):
+        return hash((self.column_name, self.index_type, self.params))
+
+    def to_ttype(self):
+        init_params_list = []
+        if self.params is not None:
+            for key, value in self.params.items():
+                if isinstance(value, str):
+                    init_params_list.append(ttypes.InitParameter(key, value))
+                else:
+                    raise InfinityException(ErrorCode.INVALID_INDEX_PARAM, f"{value} should be string type")
+
+        return ttypes.IndexInfo(
+            self.column_name.strip(),
+            self.index_type.to_ttype(),
+            init_params_list
+        )

infinity/infinity.py

@@ -0,0 +1,79 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from abc import ABC, abstractmethod
+from enum import Enum
+
+from infinity import URI
+
+
+# abstract class
+class InfinityConnection(ABC):
+    def __init__(self, uri):
+        self.uri = uri
+
+    @abstractmethod
+    def create_database(self, db_name, options=None, comment: str = None):
+        pass
+
+    @abstractmethod
+    def list_databases(self):
+        pass
+
+    @abstractmethod
+    def show_database(self, db_name):
+        pass
+
+    @abstractmethod
+    def drop_database(self, db_name, options=None):
+        pass
+
+    @abstractmethod
+    def get_database(self, db_name):
+        pass
+
+    @abstractmethod
+    def show_current_node(self):
+        pass
+
+    @abstractmethod
+    def disconnect(self):
+        pass
+
+    @abstractmethod
+    def create_database_snapshot(self, snapshot_name: str, db_name: str):
+        pass
+
+    @abstractmethod
+    def restore_database_snapshot(self, snapshot_name: str):
+        pass
+
+    @abstractmethod
+    def create_system_snapshot(self, snapshot_name: str):
+        pass
+
+    @abstractmethod
+    def restore_system_snapshot(self, snapshot_name: str):
+        pass
+
+    @abstractmethod
+    def list_snapshots(self):
+        pass
+
+    @abstractmethod
+    def show_snapshot(self, snapshot_name: str):
+        pass
+    
+    @abstractmethod
+    def drop_snapshot(self, snapshot_name: str):
+        pass