pomaidb 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pomaidb-0.1.0/PKG-INFO +62 -0
- pomaidb-0.1.0/README.md +40 -0
- pomaidb-0.1.0/pomaidb/__init__.py +657 -0
- pomaidb-0.1.0/pomaidb/lib/libpomai_c.dll +0 -0
- pomaidb-0.1.0/pomaidb/py.typed +0 -0
- pomaidb-0.1.0/pomaidb/zero_copy.py +93 -0
- pomaidb-0.1.0/pomaidb.egg-info/PKG-INFO +62 -0
- pomaidb-0.1.0/pomaidb.egg-info/SOURCES.txt +13 -0
- pomaidb-0.1.0/pomaidb.egg-info/dependency_links.txt +1 -0
- pomaidb-0.1.0/pomaidb.egg-info/not-zip-safe +1 -0
- pomaidb-0.1.0/pomaidb.egg-info/top_level.txt +1 -0
- pomaidb-0.1.0/pyproject.toml +36 -0
- pomaidb-0.1.0/setup.cfg +4 -0
- pomaidb-0.1.0/setup.py +12 -0
- pomaidb-0.1.0/tests/test_binding.py +62 -0
pomaidb-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pomaidb
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Embedded vector database for Edge AI - fast in-process C-API bindings
|
|
5
|
+
Author-email: PomaiDB Team <info@pomaidb.org>
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/pomagrenate/pomaidb
|
|
8
|
+
Project-URL: Repository, https://github.com/pomagrenate/pomaidb
|
|
9
|
+
Keywords: vector,embeddings,hnsw,database,embedded,edge-ai
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Topic :: Database
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Requires-Python: >=3.8
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# PomaiDB Python Bindings
|
|
24
|
+
|
|
25
|
+
Official Python bindings for **PomaiDB**, an embedded vector database designed for high-performance Edge AI applications.
|
|
26
|
+
|
|
27
|
+
## Installation
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
pip install pomaidb
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## Quick Start
|
|
34
|
+
|
|
35
|
+
```python
|
|
36
|
+
import pomaidb
|
|
37
|
+
|
|
38
|
+
# 1. Open database
|
|
39
|
+
db = pomaidb.open_db("test_db", dim=4)
|
|
40
|
+
|
|
41
|
+
# 2. Put vectors
|
|
42
|
+
db.put(1, [1.0, 0.0, 0.0, 0.0])
|
|
43
|
+
db.put(2, [0.0, 1.0, 0.0, 0.0], membrane="docs", payload=b"doc_payload", timestamp=1000)
|
|
44
|
+
|
|
45
|
+
# 3. Search
|
|
46
|
+
hits = db.search([1.0, 0.0, 0.0, 0.0], topk=5)
|
|
47
|
+
for hit in hits:
|
|
48
|
+
print(f"ID: {hit.id}, Score: {hit.score}")
|
|
49
|
+
|
|
50
|
+
# 4. Multi-membrane query
|
|
51
|
+
hits_docs = db.search([0.0, 1.0, 0.0, 0.0], topk=5, membrane="docs")
|
|
52
|
+
|
|
53
|
+
# 5. Flush and close
|
|
54
|
+
db.flush()
|
|
55
|
+
db.close()
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
## Features
|
|
59
|
+
- In-process embedded execution (zero network overhead, zero configuration).
|
|
60
|
+
- Pomegranate Engine Architecture: Rind MemTable, Locule immutable containers, and Press compaction.
|
|
61
|
+
- Built-in Quantization: FP32, SQ8, FP16, 1-bit binary quantization, PQ8.
|
|
62
|
+
- Native multi-membrane tenancy and arbitrary payload buffering.
|
pomaidb-0.1.0/README.md
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
# PomaiDB Python Bindings
|
|
2
|
+
|
|
3
|
+
Official Python bindings for **PomaiDB**, an embedded vector database designed for high-performance Edge AI applications.
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install pomaidb
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Quick Start
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
import pomaidb
|
|
15
|
+
|
|
16
|
+
# 1. Open database
|
|
17
|
+
db = pomaidb.open_db("test_db", dim=4)
|
|
18
|
+
|
|
19
|
+
# 2. Put vectors
|
|
20
|
+
db.put(1, [1.0, 0.0, 0.0, 0.0])
|
|
21
|
+
db.put(2, [0.0, 1.0, 0.0, 0.0], membrane="docs", payload=b"doc_payload", timestamp=1000)
|
|
22
|
+
|
|
23
|
+
# 3. Search
|
|
24
|
+
hits = db.search([1.0, 0.0, 0.0, 0.0], topk=5)
|
|
25
|
+
for hit in hits:
|
|
26
|
+
print(f"ID: {hit.id}, Score: {hit.score}")
|
|
27
|
+
|
|
28
|
+
# 4. Multi-membrane query
|
|
29
|
+
hits_docs = db.search([0.0, 1.0, 0.0, 0.0], topk=5, membrane="docs")
|
|
30
|
+
|
|
31
|
+
# 5. Flush and close
|
|
32
|
+
db.flush()
|
|
33
|
+
db.close()
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
## Features
|
|
37
|
+
- In-process embedded execution (zero network overhead, zero configuration).
|
|
38
|
+
- Pomegranate Engine Architecture: Rind MemTable, Locule immutable containers, and Press compaction.
|
|
39
|
+
- Built-in Quantization: FP32, SQ8, FP16, 1-bit binary quantization, PQ8.
|
|
40
|
+
- Native multi-membrane tenancy and arbitrary payload buffering.
|
|
@@ -0,0 +1,657 @@
|
|
|
1
|
+
"""
|
|
2
|
+
PomaiDB — embedded vector database for Edge AI.
|
|
3
|
+
|
|
4
|
+
Use the C library (libpomai_c.so / libpomai_c.dylib / pomai_c.dll) via ctypes.
|
|
5
|
+
Set POMAI_C_LIB to the path to the shared library, or build from source
|
|
6
|
+
and point to build/libpomai_c.so (Linux) or build/pomai_c.dll (Windows).
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
import ctypes
|
|
10
|
+
import json
|
|
11
|
+
import os
|
|
12
|
+
import sys
|
|
13
|
+
from pathlib import Path
|
|
14
|
+
|
|
15
|
+
__all__ = [
|
|
16
|
+
"Database", "open_db", "close", "flush", "put", "put_batch", "delete", "exists", "get",
|
|
17
|
+
"search", "search_batch", "search_zero_copy", "release_zero_copy_session",
|
|
18
|
+
"freeze", "compact", "compact_membrane", "create_membrane", "drop_membrane",
|
|
19
|
+
"open_membrane", "close_membrane", "list_membranes", "get_stats",
|
|
20
|
+
"resolve_effective_options", "PomaiDBError",
|
|
21
|
+
"MEMBRANE_KIND_VECTOR",
|
|
22
|
+
"QUANT_NONE", "QUANT_SQ8", "QUANT_FP16", "QUANT_BIT", "QUANT_PQ8",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
MEMBRANE_KIND_VECTOR = 0
|
|
26
|
+
|
|
27
|
+
QUANT_NONE = 0
|
|
28
|
+
QUANT_SQ8 = 1
|
|
29
|
+
QUANT_FP16 = 2
|
|
30
|
+
QUANT_BIT = 3
|
|
31
|
+
QUANT_PQ8 = 4
|
|
32
|
+
|
|
33
|
+
class PomaiDBError(Exception):
|
|
34
|
+
pass
|
|
35
|
+
|
|
36
|
+
def _find_lib():
|
|
37
|
+
env = os.environ.get("POMAI_C_LIB")
|
|
38
|
+
if env and os.path.exists(env):
|
|
39
|
+
return env
|
|
40
|
+
pkg_dir = Path(__file__).resolve().parent
|
|
41
|
+
for name in ["libpomai_c.dll", "pomai_c.dll", "libpomai_c.so", "libpomai_c.dylib"]:
|
|
42
|
+
if (pkg_dir / "lib" / name).is_file():
|
|
43
|
+
return str(pkg_dir / "lib" / name)
|
|
44
|
+
if (pkg_dir / name).is_file():
|
|
45
|
+
return str(pkg_dir / name)
|
|
46
|
+
for base in [Path(__file__).resolve().parents[3], Path(__file__).resolve().parents[2], Path.cwd()]:
|
|
47
|
+
for name in ["libpomai_c.so", "libpomai_c.dylib", "pomai_c.dll", "libpomai_c.dll"]:
|
|
48
|
+
p = base / "build" / name
|
|
49
|
+
if p.exists():
|
|
50
|
+
return str(p)
|
|
51
|
+
p_bin = base / "build" / "bin" / name
|
|
52
|
+
if p_bin.exists():
|
|
53
|
+
return str(p_bin)
|
|
54
|
+
return None
|
|
55
|
+
|
|
56
|
+
_lib_path = _find_lib()
|
|
57
|
+
_lib = None
|
|
58
|
+
|
|
59
|
+
def _ensure_lib():
|
|
60
|
+
global _lib
|
|
61
|
+
if _lib is not None:
|
|
62
|
+
return
|
|
63
|
+
path = _find_lib()
|
|
64
|
+
if not path or not os.path.isfile(path):
|
|
65
|
+
raise PomaiDBError(
|
|
66
|
+
"PomaiDB C library not found. Set POMAI_C_LIB to path to libpomai_c.so (or .dll/.dylib), "
|
|
67
|
+
"or build the project and run from repo root."
|
|
68
|
+
)
|
|
69
|
+
if sys.platform == "win32" and hasattr(os, "add_dll_directory"):
|
|
70
|
+
dll_dir = os.path.dirname(os.path.abspath(path))
|
|
71
|
+
if os.path.isdir(dll_dir):
|
|
72
|
+
try:
|
|
73
|
+
os.add_dll_directory(dll_dir)
|
|
74
|
+
except Exception:
|
|
75
|
+
pass
|
|
76
|
+
for p in os.environ.get("PATH", "").split(os.pathsep):
|
|
77
|
+
if p and os.path.isdir(p):
|
|
78
|
+
try:
|
|
79
|
+
os.add_dll_directory(p)
|
|
80
|
+
except Exception:
|
|
81
|
+
pass
|
|
82
|
+
_lib = ctypes.CDLL(path)
|
|
83
|
+
_register_api(_lib)
|
|
84
|
+
|
|
85
|
+
def _register_api(lib):
|
|
86
|
+
class PomaiOptions(ctypes.Structure):
|
|
87
|
+
_fields_ = [
|
|
88
|
+
("struct_size", ctypes.c_uint32),
|
|
89
|
+
("path", ctypes.c_char_p),
|
|
90
|
+
("shards", ctypes.c_uint32),
|
|
91
|
+
("dim", ctypes.c_uint32),
|
|
92
|
+
("search_threads", ctypes.c_uint32),
|
|
93
|
+
("fsync_policy", ctypes.c_int),
|
|
94
|
+
("memory_budget_bytes", ctypes.c_uint64),
|
|
95
|
+
("deadline_ms", ctypes.c_uint32),
|
|
96
|
+
("index_type", ctypes.c_uint8),
|
|
97
|
+
("hnsw_m", ctypes.c_uint32),
|
|
98
|
+
("hnsw_ef_construction", ctypes.c_uint32),
|
|
99
|
+
("hnsw_ef_search", ctypes.c_uint32),
|
|
100
|
+
("adaptive_threshold", ctypes.c_uint32),
|
|
101
|
+
("metric", ctypes.c_uint8),
|
|
102
|
+
("edge_profile", ctypes.c_uint8),
|
|
103
|
+
("tick_max_ops", ctypes.c_uint32),
|
|
104
|
+
("tick_max_ms", ctypes.c_uint32),
|
|
105
|
+
("strict_deterministic", ctypes.c_bool),
|
|
106
|
+
("quant_type", ctypes.c_uint8),
|
|
107
|
+
("pq_m", ctypes.c_uint32),
|
|
108
|
+
("memtable_flush_threshold_mb", ctypes.c_uint32),
|
|
109
|
+
("auto_freeze_on_pressure", ctypes.c_bool),
|
|
110
|
+
("max_memtable_mb", ctypes.c_uint32),
|
|
111
|
+
("write_coalesce_window_us", ctypes.c_uint32),
|
|
112
|
+
("write_coalesce_batch_size", ctypes.c_uint32),
|
|
113
|
+
("enable_encryption_at_rest", ctypes.c_bool),
|
|
114
|
+
("encryption_key_hex", ctypes.c_char_p),
|
|
115
|
+
]
|
|
116
|
+
|
|
117
|
+
class PomaiUpsert(ctypes.Structure):
|
|
118
|
+
_fields_ = [
|
|
119
|
+
("struct_size", ctypes.c_uint32),
|
|
120
|
+
("id", ctypes.c_uint64),
|
|
121
|
+
("vector", ctypes.POINTER(ctypes.c_float)),
|
|
122
|
+
("dim", ctypes.c_uint32),
|
|
123
|
+
("metadata", ctypes.POINTER(ctypes.c_uint8)),
|
|
124
|
+
("metadata_len", ctypes.c_uint32),
|
|
125
|
+
("membrane", ctypes.c_char_p),
|
|
126
|
+
("timestamp", ctypes.c_uint64),
|
|
127
|
+
("payload", ctypes.POINTER(ctypes.c_uint8)),
|
|
128
|
+
("payload_len", ctypes.c_uint32),
|
|
129
|
+
]
|
|
130
|
+
|
|
131
|
+
class PomaiQuery(ctypes.Structure):
|
|
132
|
+
_fields_ = [
|
|
133
|
+
("struct_size", ctypes.c_uint32),
|
|
134
|
+
("vector", ctypes.POINTER(ctypes.c_float)),
|
|
135
|
+
("dim", ctypes.c_uint32),
|
|
136
|
+
("topk", ctypes.c_uint32),
|
|
137
|
+
("filter_expression", ctypes.c_char_p),
|
|
138
|
+
("partition_device_id", ctypes.c_char_p),
|
|
139
|
+
("partition_location_id", ctypes.c_char_p),
|
|
140
|
+
("deadline_ms", ctypes.c_uint32),
|
|
141
|
+
("flags", ctypes.c_uint32),
|
|
142
|
+
("membrane", ctypes.c_char_p),
|
|
143
|
+
("as_of_ts", ctypes.c_uint64),
|
|
144
|
+
("as_of_lsn", ctypes.c_uint64),
|
|
145
|
+
]
|
|
146
|
+
|
|
147
|
+
class PomaiSemanticPointer(ctypes.Structure):
|
|
148
|
+
_fields_ = [
|
|
149
|
+
("struct_size", ctypes.c_uint32),
|
|
150
|
+
("raw_data_ptr", ctypes.c_void_p),
|
|
151
|
+
("dim", ctypes.c_uint32),
|
|
152
|
+
("quant_min", ctypes.c_float),
|
|
153
|
+
("quant_inv_scale", ctypes.c_float),
|
|
154
|
+
("session_id", ctypes.c_uint64),
|
|
155
|
+
]
|
|
156
|
+
|
|
157
|
+
class PomaiSearchResults(ctypes.Structure):
|
|
158
|
+
_fields_ = [
|
|
159
|
+
("struct_size", ctypes.c_uint32),
|
|
160
|
+
("count", ctypes.c_size_t),
|
|
161
|
+
("ids", ctypes.POINTER(ctypes.c_uint64)),
|
|
162
|
+
("scores", ctypes.POINTER(ctypes.c_float)),
|
|
163
|
+
("shard_ids", ctypes.POINTER(ctypes.c_uint32)),
|
|
164
|
+
("total_shards_count", ctypes.c_uint32),
|
|
165
|
+
("pruned_shards_count", ctypes.c_uint32),
|
|
166
|
+
("zero_copy_pointers", ctypes.POINTER(PomaiSemanticPointer)),
|
|
167
|
+
]
|
|
168
|
+
|
|
169
|
+
class PomaiRecord(ctypes.Structure):
|
|
170
|
+
_fields_ = [
|
|
171
|
+
("struct_size", ctypes.c_uint32),
|
|
172
|
+
("id", ctypes.c_uint64),
|
|
173
|
+
("dim", ctypes.c_uint32),
|
|
174
|
+
("vector", ctypes.POINTER(ctypes.c_float)),
|
|
175
|
+
("metadata", ctypes.POINTER(ctypes.c_uint8)),
|
|
176
|
+
("metadata_len", ctypes.c_uint32),
|
|
177
|
+
("is_deleted", ctypes.c_bool),
|
|
178
|
+
("timestamp", ctypes.c_uint64),
|
|
179
|
+
("payload", ctypes.POINTER(ctypes.c_uint8)),
|
|
180
|
+
("payload_len", ctypes.c_uint32),
|
|
181
|
+
]
|
|
182
|
+
|
|
183
|
+
lib.PomaiOptions = PomaiOptions
|
|
184
|
+
lib.PomaiUpsert = PomaiUpsert
|
|
185
|
+
lib.PomaiQuery = PomaiQuery
|
|
186
|
+
lib.PomaiSemanticPointer = PomaiSemanticPointer
|
|
187
|
+
lib.PomaiSearchResults = PomaiSearchResults
|
|
188
|
+
lib.PomaiRecord = PomaiRecord
|
|
189
|
+
|
|
190
|
+
lib.pomai_options_init.argtypes = [ctypes.POINTER(PomaiOptions)]
|
|
191
|
+
lib.pomai_options_init.restype = None
|
|
192
|
+
|
|
193
|
+
lib.pomai_open.argtypes = [ctypes.POINTER(PomaiOptions), ctypes.POINTER(ctypes.c_void_p)]
|
|
194
|
+
lib.pomai_open.restype = ctypes.c_void_p
|
|
195
|
+
|
|
196
|
+
lib.pomai_close.argtypes = [ctypes.c_void_p]
|
|
197
|
+
lib.pomai_close.restype = ctypes.c_void_p
|
|
198
|
+
|
|
199
|
+
lib.pomai_flush.argtypes = [ctypes.c_void_p]
|
|
200
|
+
lib.pomai_flush.restype = ctypes.c_void_p
|
|
201
|
+
|
|
202
|
+
lib.pomai_compact.argtypes = [ctypes.c_void_p]
|
|
203
|
+
lib.pomai_compact.restype = ctypes.c_void_p
|
|
204
|
+
|
|
205
|
+
lib.pomai_freeze.argtypes = [ctypes.c_void_p]
|
|
206
|
+
lib.pomai_freeze.restype = ctypes.c_void_p
|
|
207
|
+
|
|
208
|
+
lib.pomai_freeze_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
|
|
209
|
+
lib.pomai_freeze_membrane.restype = ctypes.c_void_p
|
|
210
|
+
|
|
211
|
+
lib.pomai_put.argtypes = [ctypes.c_void_p, ctypes.POINTER(PomaiUpsert)]
|
|
212
|
+
lib.pomai_put.restype = ctypes.c_void_p
|
|
213
|
+
|
|
214
|
+
lib.pomai_put_batch.argtypes = [ctypes.c_void_p, ctypes.POINTER(PomaiUpsert), ctypes.c_size_t]
|
|
215
|
+
lib.pomai_put_batch.restype = ctypes.c_void_p
|
|
216
|
+
|
|
217
|
+
lib.pomai_delete.argtypes = [ctypes.c_void_p, ctypes.c_uint64]
|
|
218
|
+
lib.pomai_delete.restype = ctypes.c_void_p
|
|
219
|
+
|
|
220
|
+
lib.pomai_exists.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_bool)]
|
|
221
|
+
lib.pomai_exists.restype = ctypes.c_void_p
|
|
222
|
+
|
|
223
|
+
lib.pomai_get.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(PomaiRecord))]
|
|
224
|
+
lib.pomai_get.restype = ctypes.c_void_p
|
|
225
|
+
|
|
226
|
+
lib.pomai_record_free.argtypes = [ctypes.POINTER(PomaiRecord)]
|
|
227
|
+
lib.pomai_record_free.restype = None
|
|
228
|
+
|
|
229
|
+
lib.pomai_search.argtypes = [ctypes.c_void_p, ctypes.POINTER(PomaiQuery), ctypes.POINTER(ctypes.POINTER(PomaiSearchResults))]
|
|
230
|
+
lib.pomai_search.restype = ctypes.c_void_p
|
|
231
|
+
|
|
232
|
+
lib.pomai_search_results_free.argtypes = [ctypes.POINTER(PomaiSearchResults)]
|
|
233
|
+
lib.pomai_search_results_free.restype = None
|
|
234
|
+
|
|
235
|
+
lib.pomai_search_batch.argtypes = [
|
|
236
|
+
ctypes.c_void_p, ctypes.POINTER(PomaiQuery), ctypes.c_size_t,
|
|
237
|
+
ctypes.POINTER(ctypes.POINTER(PomaiSearchResults))
|
|
238
|
+
]
|
|
239
|
+
lib.pomai_search_batch.restype = ctypes.c_void_p
|
|
240
|
+
|
|
241
|
+
lib.pomai_search_batch_free.argtypes = [ctypes.POINTER(PomaiSearchResults), ctypes.c_size_t]
|
|
242
|
+
lib.pomai_search_batch_free.restype = None
|
|
243
|
+
|
|
244
|
+
lib.pomai_create_membrane_kind.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.uint32 if hasattr(ctypes, "uint32") else ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32]
|
|
245
|
+
lib.pomai_create_membrane_kind.restype = ctypes.c_void_p
|
|
246
|
+
|
|
247
|
+
lib.pomai_drop_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
|
|
248
|
+
lib.pomai_drop_membrane.restype = ctypes.c_void_p
|
|
249
|
+
|
|
250
|
+
lib.pomai_open_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
|
|
251
|
+
lib.pomai_open_membrane.restype = ctypes.c_void_p
|
|
252
|
+
|
|
253
|
+
lib.pomai_close_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
|
|
254
|
+
lib.pomai_close_membrane.restype = ctypes.c_void_p
|
|
255
|
+
|
|
256
|
+
lib.pomai_list_membranes_json.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_char_p), ctypes.POINTER(ctypes.c_size_t)]
|
|
257
|
+
lib.pomai_list_membranes_json.restype = ctypes.c_void_p
|
|
258
|
+
|
|
259
|
+
lib.pomai_compact_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
|
|
260
|
+
lib.pomai_compact_membrane.restype = ctypes.c_void_p
|
|
261
|
+
|
|
262
|
+
lib.pomai_get_stats_json.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_char_p), ctypes.POINTER(ctypes.c_size_t)]
|
|
263
|
+
lib.pomai_get_stats_json.restype = ctypes.c_void_p
|
|
264
|
+
|
|
265
|
+
lib.pomai_put_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(PomaiUpsert)]
|
|
266
|
+
lib.pomai_put_membrane.restype = ctypes.c_void_p
|
|
267
|
+
|
|
268
|
+
lib.pomai_put_batch_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(PomaiUpsert), ctypes.c_size_t]
|
|
269
|
+
lib.pomai_put_batch_membrane.restype = ctypes.c_void_p
|
|
270
|
+
|
|
271
|
+
lib.pomai_get_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(PomaiRecord))]
|
|
272
|
+
lib.pomai_get_membrane.restype = ctypes.c_void_p
|
|
273
|
+
|
|
274
|
+
lib.pomai_delete_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint64]
|
|
275
|
+
lib.pomai_delete_membrane.restype = ctypes.c_void_p
|
|
276
|
+
|
|
277
|
+
lib.pomai_exists_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_bool)]
|
|
278
|
+
lib.pomai_exists_membrane.restype = ctypes.c_void_p
|
|
279
|
+
|
|
280
|
+
lib.pomai_search_membrane.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(PomaiQuery), ctypes.POINTER(ctypes.POINTER(PomaiSearchResults))]
|
|
281
|
+
lib.pomai_search_membrane.restype = ctypes.c_void_p
|
|
282
|
+
|
|
283
|
+
lib.pomai_release_pointer.argtypes = [ctypes.c_uint64]
|
|
284
|
+
lib.pomai_release_pointer.restype = None
|
|
285
|
+
|
|
286
|
+
lib.pomai_free.argtypes = [ctypes.c_void_p]
|
|
287
|
+
lib.pomai_free.restype = None
|
|
288
|
+
|
|
289
|
+
lib.pomai_status_message.argtypes = [ctypes.c_void_p]
|
|
290
|
+
lib.pomai_status_message.restype = ctypes.c_char_p
|
|
291
|
+
|
|
292
|
+
lib.pomai_status_free.argtypes = [ctypes.c_void_p]
|
|
293
|
+
lib.pomai_status_free.restype = None
|
|
294
|
+
|
|
295
|
+
def _check_status(st):
|
|
296
|
+
if not st:
|
|
297
|
+
return
|
|
298
|
+
msg = _lib.pomai_status_message(st)
|
|
299
|
+
err = msg.decode("utf-8", errors="replace") if msg else "Unknown PomaiDB error"
|
|
300
|
+
_lib.pomai_status_free(st)
|
|
301
|
+
raise PomaiDBError(err)
|
|
302
|
+
|
|
303
|
+
def open_db(path, dim, shards=1, metric="l2", edge_profile=0, quant_type=0, memory_budget_bytes=0, auto_freeze_on_pressure=True, memtable_flush_threshold_mb=64):
|
|
304
|
+
_ensure_lib()
|
|
305
|
+
opts = _lib.PomaiOptions()
|
|
306
|
+
_lib.pomai_options_init(ctypes.byref(opts))
|
|
307
|
+
opts.path = path.encode("utf-8")
|
|
308
|
+
opts.dim = dim
|
|
309
|
+
opts.shards = shards
|
|
310
|
+
opts.metric = 2 if metric.lower() == "cosine" else (1 if metric.lower() in ("ip", "innerproduct") else 0)
|
|
311
|
+
opts.edge_profile = edge_profile
|
|
312
|
+
opts.quant_type = quant_type
|
|
313
|
+
opts.memory_budget_bytes = memory_budget_bytes
|
|
314
|
+
opts.auto_freeze_on_pressure = auto_freeze_on_pressure
|
|
315
|
+
opts.memtable_flush_threshold_mb = memtable_flush_threshold_mb
|
|
316
|
+
db_ptr = ctypes.c_void_p()
|
|
317
|
+
st = _lib.pomai_open(ctypes.byref(opts), ctypes.byref(db_ptr))
|
|
318
|
+
_check_status(st)
|
|
319
|
+
return db_ptr
|
|
320
|
+
|
|
321
|
+
def close(db):
|
|
322
|
+
if db:
|
|
323
|
+
_check_status(_lib.pomai_close(db))
|
|
324
|
+
|
|
325
|
+
def flush(db):
|
|
326
|
+
if db:
|
|
327
|
+
_check_status(_lib.pomai_flush(db))
|
|
328
|
+
|
|
329
|
+
def freeze(db, membrane=""):
|
|
330
|
+
if membrane:
|
|
331
|
+
_check_status(_lib.pomai_freeze_membrane(db, membrane.encode("utf-8")))
|
|
332
|
+
else:
|
|
333
|
+
_check_status(_lib.pomai_freeze(db))
|
|
334
|
+
|
|
335
|
+
def compact(db):
|
|
336
|
+
if db:
|
|
337
|
+
_check_status(_lib.pomai_compact(db))
|
|
338
|
+
|
|
339
|
+
def compact_membrane(db, name):
|
|
340
|
+
_check_status(_lib.pomai_compact_membrane(db, name.encode("utf-8")))
|
|
341
|
+
|
|
342
|
+
def put(db, id, vector, tenant="", membrane="", timestamp=0, payload=b""):
|
|
343
|
+
_ensure_lib()
|
|
344
|
+
up = _lib.PomaiUpsert()
|
|
345
|
+
up.struct_size = ctypes.sizeof(_lib.PomaiUpsert)
|
|
346
|
+
up.id = id
|
|
347
|
+
up.dim = len(vector)
|
|
348
|
+
c_floats = (ctypes.c_float * len(vector))(*vector)
|
|
349
|
+
up.vector = c_floats
|
|
350
|
+
if tenant:
|
|
351
|
+
t_bytes = tenant.encode("utf-8")
|
|
352
|
+
up.metadata = (ctypes.c_uint8 * len(t_bytes))(*t_bytes)
|
|
353
|
+
up.metadata_len = len(t_bytes)
|
|
354
|
+
if membrane:
|
|
355
|
+
up.membrane = membrane.encode("utf-8")
|
|
356
|
+
up.timestamp = timestamp
|
|
357
|
+
if payload:
|
|
358
|
+
p_bytes = payload if isinstance(payload, (bytes, bytearray)) else str(payload).encode("utf-8")
|
|
359
|
+
up.payload = (ctypes.c_uint8 * len(p_bytes))(*p_bytes)
|
|
360
|
+
up.payload_len = len(p_bytes)
|
|
361
|
+
_check_status(_lib.pomai_put(db, ctypes.byref(up)))
|
|
362
|
+
|
|
363
|
+
def put_batch(db, ids, vectors, tenants=None, membrane=""):
|
|
364
|
+
_ensure_lib()
|
|
365
|
+
n = len(ids)
|
|
366
|
+
if n == 0:
|
|
367
|
+
return
|
|
368
|
+
arr_type = _lib.PomaiUpsert * n
|
|
369
|
+
arr = arr_type()
|
|
370
|
+
keep_alive = []
|
|
371
|
+
memb_bytes = membrane.encode("utf-8") if membrane else None
|
|
372
|
+
for i in range(n):
|
|
373
|
+
arr[i].struct_size = ctypes.sizeof(_lib.PomaiUpsert)
|
|
374
|
+
arr[i].id = ids[i]
|
|
375
|
+
arr[i].dim = len(vectors[i])
|
|
376
|
+
c_v = (ctypes.c_float * len(vectors[i]))(*vectors[i])
|
|
377
|
+
arr[i].vector = c_v
|
|
378
|
+
keep_alive.append(c_v)
|
|
379
|
+
if memb_bytes:
|
|
380
|
+
arr[i].membrane = memb_bytes
|
|
381
|
+
if tenants and i < len(tenants) and tenants[i]:
|
|
382
|
+
t_b = tenants[i].encode("utf-8")
|
|
383
|
+
c_m = (ctypes.c_uint8 * len(t_b))(*t_b)
|
|
384
|
+
arr[i].metadata = c_m
|
|
385
|
+
arr[i].metadata_len = len(t_b)
|
|
386
|
+
keep_alive.append(c_m)
|
|
387
|
+
_check_status(_lib.pomai_put_batch(db, arr, n))
|
|
388
|
+
|
|
389
|
+
def delete(db, id, membrane=""):
|
|
390
|
+
if membrane:
|
|
391
|
+
_check_status(_lib.pomai_delete_membrane(db, membrane.encode("utf-8"), id))
|
|
392
|
+
else:
|
|
393
|
+
_check_status(_lib.pomai_delete(db, id))
|
|
394
|
+
|
|
395
|
+
def exists(db, id, membrane=""):
|
|
396
|
+
out = ctypes.c_bool()
|
|
397
|
+
if membrane:
|
|
398
|
+
_check_status(_lib.pomai_exists_membrane(db, membrane.encode("utf-8"), id, ctypes.byref(out)))
|
|
399
|
+
else:
|
|
400
|
+
_check_status(_lib.pomai_exists(db, id, ctypes.byref(out)))
|
|
401
|
+
return out.value
|
|
402
|
+
|
|
403
|
+
class Hit(tuple):
|
|
404
|
+
"""Search hit (id, score) supporting tuple indexing, dict keys, and attributes."""
|
|
405
|
+
def __new__(cls, id, score):
|
|
406
|
+
return super(Hit, cls).__new__(cls, (id, score))
|
|
407
|
+
@property
|
|
408
|
+
def id(self):
|
|
409
|
+
return self[0]
|
|
410
|
+
@property
|
|
411
|
+
def score(self):
|
|
412
|
+
return self[1]
|
|
413
|
+
def __getitem__(self, item):
|
|
414
|
+
if item == "id":
|
|
415
|
+
return self[0]
|
|
416
|
+
if item == "score":
|
|
417
|
+
return self[1]
|
|
418
|
+
return super().__getitem__(item)
|
|
419
|
+
def __repr__(self):
|
|
420
|
+
return f"Hit(id={self[0]}, score={self[1]})"
|
|
421
|
+
|
|
422
|
+
def get(db, id, membrane=""):
|
|
423
|
+
rec_ptr = ctypes.POINTER(_lib.PomaiRecord)()
|
|
424
|
+
if membrane:
|
|
425
|
+
_check_status(_lib.pomai_get_membrane(db, membrane.encode("utf-8"), id, ctypes.byref(rec_ptr)))
|
|
426
|
+
else:
|
|
427
|
+
_check_status(_lib.pomai_get(db, id, ctypes.byref(rec_ptr)))
|
|
428
|
+
if not rec_ptr:
|
|
429
|
+
return None
|
|
430
|
+
r = rec_ptr.contents
|
|
431
|
+
vec = [r.vector[i] for i in range(r.dim)]
|
|
432
|
+
meta = ""
|
|
433
|
+
if r.metadata and r.metadata_len > 0:
|
|
434
|
+
meta = bytes(r.metadata[:r.metadata_len]).decode("utf-8", errors="replace")
|
|
435
|
+
payload = b""
|
|
436
|
+
if r.payload and r.payload_len > 0:
|
|
437
|
+
payload = bytes(r.payload[:r.payload_len])
|
|
438
|
+
dim = r.dim
|
|
439
|
+
ts = r.timestamp
|
|
440
|
+
_lib.pomai_record_free(rec_ptr)
|
|
441
|
+
return {"id": id, "dim": dim, "vector": vec, "tenant": meta, "timestamp": ts, "payload": payload}
|
|
442
|
+
|
|
443
|
+
def search(db, query_vector, topk=10, tenant="", membrane="", as_of_ts=0, as_of_lsn=0):
|
|
444
|
+
_ensure_lib()
|
|
445
|
+
q = _lib.PomaiQuery()
|
|
446
|
+
q.struct_size = ctypes.sizeof(_lib.PomaiQuery)
|
|
447
|
+
c_v = (ctypes.c_float * len(query_vector))(*query_vector)
|
|
448
|
+
q.vector = c_v
|
|
449
|
+
q.dim = len(query_vector)
|
|
450
|
+
q.topk = topk
|
|
451
|
+
if tenant:
|
|
452
|
+
q.filter_expression = f"tenant={tenant}".encode("utf-8")
|
|
453
|
+
if membrane:
|
|
454
|
+
q.membrane = membrane.encode("utf-8")
|
|
455
|
+
q.as_of_ts = as_of_ts
|
|
456
|
+
q.as_of_lsn = as_of_lsn
|
|
457
|
+
res_ptr = ctypes.POINTER(_lib.PomaiSearchResults)()
|
|
458
|
+
_check_status(_lib.pomai_search(db, ctypes.byref(q), ctypes.byref(res_ptr)))
|
|
459
|
+
if not res_ptr:
|
|
460
|
+
return []
|
|
461
|
+
res = res_ptr.contents
|
|
462
|
+
hits = [Hit(res.ids[i], res.scores[i]) for i in range(res.count)]
|
|
463
|
+
_lib.pomai_search_results_free(res_ptr)
|
|
464
|
+
return hits
|
|
465
|
+
|
|
466
|
+
def search_batch(db, query_vectors, topk=10, membrane=""):
|
|
467
|
+
_ensure_lib()
|
|
468
|
+
n = len(query_vectors)
|
|
469
|
+
if n == 0:
|
|
470
|
+
return []
|
|
471
|
+
arr_type = _lib.PomaiQuery * n
|
|
472
|
+
arr = arr_type()
|
|
473
|
+
keep_alive = []
|
|
474
|
+
dim = len(query_vectors[0])
|
|
475
|
+
memb_bytes = membrane.encode("utf-8") if membrane else None
|
|
476
|
+
for i in range(n):
|
|
477
|
+
arr[i].struct_size = ctypes.sizeof(_lib.PomaiQuery)
|
|
478
|
+
c_v = (ctypes.c_float * dim)(*query_vectors[i])
|
|
479
|
+
arr[i].vector = c_v
|
|
480
|
+
arr[i].dim = dim
|
|
481
|
+
arr[i].topk = topk
|
|
482
|
+
if memb_bytes:
|
|
483
|
+
arr[i].membrane = memb_bytes
|
|
484
|
+
keep_alive.append(c_v)
|
|
485
|
+
res_ptr = ctypes.POINTER(_lib.PomaiSearchResults)()
|
|
486
|
+
_check_status(_lib.pomai_search_batch(db, arr, n, ctypes.byref(res_ptr)))
|
|
487
|
+
if not res_ptr:
|
|
488
|
+
return []
|
|
489
|
+
out = []
|
|
490
|
+
for i in range(n):
|
|
491
|
+
r = res_ptr[i]
|
|
492
|
+
hits = [Hit(r.ids[j], r.scores[j]) for j in range(r.count)] if r.count > 0 else []
|
|
493
|
+
out.append(hits)
|
|
494
|
+
_lib.pomai_search_batch_free(res_ptr, n)
|
|
495
|
+
return out
|
|
496
|
+
|
|
497
|
+
def search_zero_copy(db, query_vector, topk=10, membrane=""):
|
|
498
|
+
_ensure_lib()
|
|
499
|
+
q = _lib.PomaiQuery()
|
|
500
|
+
q.struct_size = ctypes.sizeof(_lib.PomaiQuery)
|
|
501
|
+
c_v = (ctypes.c_float * len(query_vector))(*query_vector)
|
|
502
|
+
q.vector = c_v
|
|
503
|
+
q.dim = len(query_vector)
|
|
504
|
+
q.topk = topk
|
|
505
|
+
q.flags = 1 # ZERO_COPY
|
|
506
|
+
if membrane:
|
|
507
|
+
q.membrane = membrane.encode("utf-8")
|
|
508
|
+
res_ptr = ctypes.POINTER(_lib.PomaiSearchResults)()
|
|
509
|
+
_check_status(_lib.pomai_search(db, ctypes.byref(q), ctypes.byref(res_ptr)))
|
|
510
|
+
if not res_ptr:
|
|
511
|
+
return {"hits": [], "session_id": 0}
|
|
512
|
+
res = res_ptr.contents
|
|
513
|
+
hits = [{"id": res.ids[i], "score": res.scores[i]} for i in range(res.count)]
|
|
514
|
+
sess_id = res.zero_copy_pointers[0].session_id if res.count > 0 and res.zero_copy_pointers else 0
|
|
515
|
+
_lib.pomai_search_results_free(res_ptr)
|
|
516
|
+
return {"hits": hits, "session_id": sess_id}
|
|
517
|
+
|
|
518
|
+
def release_zero_copy_session(session_id):
|
|
519
|
+
if session_id:
|
|
520
|
+
_lib.pomai_release_pointer(session_id)
|
|
521
|
+
|
|
522
|
+
def create_membrane(db, name, dim, shard_count=1):
|
|
523
|
+
_check_status(_lib.pomai_create_membrane_kind(db, name.encode("utf-8"), dim, shard_count, 0))
|
|
524
|
+
|
|
525
|
+
def drop_membrane(db, name):
|
|
526
|
+
_check_status(_lib.pomai_drop_membrane(db, name.encode("utf-8")))
|
|
527
|
+
|
|
528
|
+
def open_membrane(db, name):
|
|
529
|
+
_check_status(_lib.pomai_open_membrane(db, name.encode("utf-8")))
|
|
530
|
+
|
|
531
|
+
def close_membrane(db, name):
|
|
532
|
+
_check_status(_lib.pomai_close_membrane(db, name.encode("utf-8")))
|
|
533
|
+
|
|
534
|
+
def list_membranes(db):
|
|
535
|
+
out_json = ctypes.c_char_p()
|
|
536
|
+
out_len = ctypes.c_size_t()
|
|
537
|
+
_check_status(_lib.pomai_list_membranes_json(db, ctypes.byref(out_json), ctypes.byref(out_len)))
|
|
538
|
+
if not out_json.value:
|
|
539
|
+
return []
|
|
540
|
+
s = out_json.value.decode("utf-8")
|
|
541
|
+
_lib.pomai_free(ctypes.cast(out_json, ctypes.c_void_p))
|
|
542
|
+
return json.loads(s)
|
|
543
|
+
|
|
544
|
+
def get_stats(db):
|
|
545
|
+
out_json = ctypes.c_char_p()
|
|
546
|
+
out_len = ctypes.c_size_t()
|
|
547
|
+
_check_status(_lib.pomai_get_stats_json(db, ctypes.byref(out_json), ctypes.byref(out_len)))
|
|
548
|
+
if not out_json.value:
|
|
549
|
+
return {}
|
|
550
|
+
s = out_json.value.decode("utf-8")
|
|
551
|
+
_lib.pomai_free(ctypes.cast(out_json, ctypes.c_void_p))
|
|
552
|
+
return json.loads(s)
|
|
553
|
+
|
|
554
|
+
def resolve_effective_options(path, dim, shards=1, edge_profile=0):
|
|
555
|
+
_ensure_lib()
|
|
556
|
+
opts = _lib.PomaiOptions()
|
|
557
|
+
_lib.pomai_options_init(ctypes.byref(opts))
|
|
558
|
+
opts.path = path.encode("utf-8")
|
|
559
|
+
opts.dim = dim
|
|
560
|
+
opts.shards = shards
|
|
561
|
+
opts.edge_profile = edge_profile
|
|
562
|
+
out_json = ctypes.c_char_p()
|
|
563
|
+
out_len = ctypes.c_size_t()
|
|
564
|
+
_check_status(_lib.pomai_options_resolve_json(ctypes.byref(opts), ctypes.byref(out_json), ctypes.byref(out_len)))
|
|
565
|
+
s = out_json.value.decode("utf-8")
|
|
566
|
+
_lib.pomai_free(ctypes.cast(out_json, ctypes.c_void_p))
|
|
567
|
+
return json.loads(s)
|
|
568
|
+
|
|
569
|
+
|
|
570
|
+
class Database:
|
|
571
|
+
"""Object-oriented wrapper around a PomaiDB database handle."""
|
|
572
|
+
def __init__(self, handle):
|
|
573
|
+
self._handle = handle
|
|
574
|
+
|
|
575
|
+
@classmethod
|
|
576
|
+
def open(cls, path, dim, shards=1, metric="l2", quant_type=QUANT_NONE, **kwargs):
|
|
577
|
+
h = open_db(path, dim, shards=shards, metric=metric, quant_type=quant_type, **kwargs)
|
|
578
|
+
return cls(h)
|
|
579
|
+
|
|
580
|
+
def close(self):
|
|
581
|
+
if self._handle:
|
|
582
|
+
close(self._handle)
|
|
583
|
+
self._handle = None
|
|
584
|
+
|
|
585
|
+
def __enter__(self):
|
|
586
|
+
return self
|
|
587
|
+
|
|
588
|
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
589
|
+
self.close()
|
|
590
|
+
|
|
591
|
+
def put(self, id, vector, membrane=None, timestamp=0, payload=None):
|
|
592
|
+
put(self._handle, id, vector, membrane=membrane, timestamp=timestamp, payload=payload)
|
|
593
|
+
|
|
594
|
+
def put_batch(self, ids, vectors, membrane=None):
|
|
595
|
+
put_batch(self._handle, ids, vectors, membrane=membrane)
|
|
596
|
+
|
|
597
|
+
def get(self, id, membrane=None, with_metadata=False):
|
|
598
|
+
res = get(self._handle, id, membrane=membrane)
|
|
599
|
+
if not with_metadata and isinstance(res, dict) and "vector" in res:
|
|
600
|
+
return res["vector"]
|
|
601
|
+
if with_metadata and isinstance(res, dict):
|
|
602
|
+
class Record:
|
|
603
|
+
def __init__(self, d):
|
|
604
|
+
self.id = d.get("id")
|
|
605
|
+
self.vector = d.get("vector")
|
|
606
|
+
self.dim = d.get("dim")
|
|
607
|
+
self.timestamp = d.get("timestamp")
|
|
608
|
+
self.payload = d.get("payload")
|
|
609
|
+
return Record(res)
|
|
610
|
+
return res
|
|
611
|
+
|
|
612
|
+
def exists(self, id, membrane=None):
|
|
613
|
+
return exists(self._handle, id, membrane=membrane)
|
|
614
|
+
|
|
615
|
+
def delete(self, id, membrane=None):
|
|
616
|
+
delete(self._handle, id, membrane=membrane)
|
|
617
|
+
|
|
618
|
+
def search(self, query, topk=10, membrane=None, as_of_ts=0):
|
|
619
|
+
class Hit:
|
|
620
|
+
def __init__(self, item):
|
|
621
|
+
self.id = item[0]
|
|
622
|
+
self.score = item[1]
|
|
623
|
+
def __repr__(self):
|
|
624
|
+
return f"Hit(id={self.id}, score={self.score})"
|
|
625
|
+
raw = search(self._handle, query, topk=topk, membrane=membrane, as_of_ts=as_of_ts)
|
|
626
|
+
return [Hit(item) for item in raw]
|
|
627
|
+
|
|
628
|
+
def search_batch(self, queries, topk=10, membrane=None, as_of_ts=0):
|
|
629
|
+
return search_batch(self._handle, queries, topk=topk, membrane=membrane, as_of_ts=as_of_ts)
|
|
630
|
+
|
|
631
|
+
def flush(self):
|
|
632
|
+
flush(self._handle)
|
|
633
|
+
|
|
634
|
+
def freeze(self, membrane=None):
|
|
635
|
+
freeze(self._handle, membrane=membrane)
|
|
636
|
+
|
|
637
|
+
def compact(self, membrane=None):
|
|
638
|
+
compact(self._handle, membrane=membrane)
|
|
639
|
+
|
|
640
|
+
def create_membrane(self, name, dim, shard_count=1):
|
|
641
|
+
create_membrane(self._handle, name, dim, shard_count=shard_count)
|
|
642
|
+
|
|
643
|
+
def drop_membrane(self, name):
|
|
644
|
+
drop_membrane(self._handle, name)
|
|
645
|
+
|
|
646
|
+
def open_membrane(self, name):
|
|
647
|
+
open_membrane(self._handle, name)
|
|
648
|
+
|
|
649
|
+
def close_membrane(self, name):
|
|
650
|
+
close_membrane(self._handle, name)
|
|
651
|
+
|
|
652
|
+
def list_membranes(self):
|
|
653
|
+
return list_membranes(self._handle)
|
|
654
|
+
|
|
655
|
+
def get_stats(self):
|
|
656
|
+
return get_stats(self._handle)
|
|
657
|
+
|
|
Binary file
|
|
File without changes
|
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Zero-copy helpers for PomaiDB search results.
|
|
3
|
+
|
|
4
|
+
This module uses the existing C ABI zero-copy semantic pointer path
|
|
5
|
+
(`POMAI_QUERY_FLAG_ZERO_COPY`) and maps the returned memory into NumPy
|
|
6
|
+
without copying.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import ctypes
|
|
12
|
+
|
|
13
|
+
from . import _ensure_lib, _lib, PomaiDBError
|
|
14
|
+
|
|
15
|
+
POMAI_QUERY_FLAG_ZERO_COPY = 1
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def release_zero_copy_session(session_id: int) -> None:
|
|
19
|
+
"""Release a pinned zero-copy session id."""
|
|
20
|
+
_ensure_lib()
|
|
21
|
+
_lib.pomai_release_pointer(int(session_id))
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def search_zero_copy(db, query, topk: int = 10):
|
|
25
|
+
"""
|
|
26
|
+
Execute one query with zero-copy enabled.
|
|
27
|
+
|
|
28
|
+
Returns list of dicts:
|
|
29
|
+
{"id", "score", "raw_u8", "dequant_f32", "session_id"}
|
|
30
|
+
"""
|
|
31
|
+
_ensure_lib()
|
|
32
|
+
try:
|
|
33
|
+
import numpy as np
|
|
34
|
+
except Exception as exc:
|
|
35
|
+
raise PomaiDBError("numpy is required for zero-copy helpers") from exc
|
|
36
|
+
|
|
37
|
+
dim = len(query)
|
|
38
|
+
qvec = (ctypes.c_float * dim)(*query)
|
|
39
|
+
q = _lib._pomai_query()
|
|
40
|
+
q.struct_size = ctypes.sizeof(_lib._pomai_query())
|
|
41
|
+
q.vector = qvec
|
|
42
|
+
q.dim = dim
|
|
43
|
+
q.topk = int(topk)
|
|
44
|
+
q.filter_expression = None
|
|
45
|
+
q.partition_device_id = None
|
|
46
|
+
q.partition_location_id = None
|
|
47
|
+
q.as_of_ts = 0
|
|
48
|
+
q.as_of_lsn = 0
|
|
49
|
+
q.aggregate_op = 0
|
|
50
|
+
q.aggregate_field = None
|
|
51
|
+
q.aggregate_topk = 0
|
|
52
|
+
q.mesh_detail_preference = 0
|
|
53
|
+
q.alpha = 1.0
|
|
54
|
+
q.deadline_ms = 0
|
|
55
|
+
q.flags = POMAI_QUERY_FLAG_ZERO_COPY
|
|
56
|
+
|
|
57
|
+
out = ctypes.POINTER(_lib._pomai_search_results)()
|
|
58
|
+
st = _lib.pomai_search(db, ctypes.byref(q), ctypes.byref(out))
|
|
59
|
+
if st:
|
|
60
|
+
msg = _lib.pomai_status_message(st).decode("utf-8", errors="replace")
|
|
61
|
+
_lib.pomai_status_free(st)
|
|
62
|
+
raise PomaiDBError(msg)
|
|
63
|
+
|
|
64
|
+
rows = []
|
|
65
|
+
try:
|
|
66
|
+
count = out.contents.count
|
|
67
|
+
pointers = out.contents.zero_copy_pointers
|
|
68
|
+
for i in range(min(int(topk), count)):
|
|
69
|
+
session_id = 0
|
|
70
|
+
raw_u8 = None
|
|
71
|
+
dequant_f32 = None
|
|
72
|
+
if pointers:
|
|
73
|
+
p = pointers[i]
|
|
74
|
+
session_id = int(p.session_id)
|
|
75
|
+
if p.raw_data_ptr and p.dim > 0:
|
|
76
|
+
addr = int(ctypes.cast(p.raw_data_ptr, ctypes.c_void_p).value or 0)
|
|
77
|
+
if addr:
|
|
78
|
+
buf = (ctypes.c_uint8 * int(p.dim)).from_address(addr)
|
|
79
|
+
raw_u8 = np.frombuffer(buf, dtype=np.uint8, count=int(p.dim))
|
|
80
|
+
dequant_f32 = raw_u8.astype(np.float32) * float(p.quant_inv_scale) + float(p.quant_min)
|
|
81
|
+
rows.append(
|
|
82
|
+
{
|
|
83
|
+
"id": int(out.contents.ids[i]),
|
|
84
|
+
"score": float(out.contents.scores[i]),
|
|
85
|
+
"raw_u8": raw_u8,
|
|
86
|
+
"dequant_f32": dequant_f32,
|
|
87
|
+
"session_id": session_id,
|
|
88
|
+
}
|
|
89
|
+
)
|
|
90
|
+
return rows
|
|
91
|
+
finally:
|
|
92
|
+
_lib.pomai_search_results_free(out)
|
|
93
|
+
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pomaidb
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Embedded vector database for Edge AI - fast in-process C-API bindings
|
|
5
|
+
Author-email: PomaiDB Team <info@pomaidb.org>
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/pomagrenate/pomaidb
|
|
8
|
+
Project-URL: Repository, https://github.com/pomagrenate/pomaidb
|
|
9
|
+
Keywords: vector,embeddings,hnsw,database,embedded,edge-ai
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Topic :: Database
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Requires-Python: >=3.8
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# PomaiDB Python Bindings
|
|
24
|
+
|
|
25
|
+
Official Python bindings for **PomaiDB**, an embedded vector database designed for high-performance Edge AI applications.
|
|
26
|
+
|
|
27
|
+
## Installation
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
pip install pomaidb
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## Quick Start
|
|
34
|
+
|
|
35
|
+
```python
|
|
36
|
+
import pomaidb
|
|
37
|
+
|
|
38
|
+
# 1. Open database
|
|
39
|
+
db = pomaidb.open_db("test_db", dim=4)
|
|
40
|
+
|
|
41
|
+
# 2. Put vectors
|
|
42
|
+
db.put(1, [1.0, 0.0, 0.0, 0.0])
|
|
43
|
+
db.put(2, [0.0, 1.0, 0.0, 0.0], membrane="docs", payload=b"doc_payload", timestamp=1000)
|
|
44
|
+
|
|
45
|
+
# 3. Search
|
|
46
|
+
hits = db.search([1.0, 0.0, 0.0, 0.0], topk=5)
|
|
47
|
+
for hit in hits:
|
|
48
|
+
print(f"ID: {hit.id}, Score: {hit.score}")
|
|
49
|
+
|
|
50
|
+
# 4. Multi-membrane query
|
|
51
|
+
hits_docs = db.search([0.0, 1.0, 0.0, 0.0], topk=5, membrane="docs")
|
|
52
|
+
|
|
53
|
+
# 5. Flush and close
|
|
54
|
+
db.flush()
|
|
55
|
+
db.close()
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
## Features
|
|
59
|
+
- In-process embedded execution (zero network overhead, zero configuration).
|
|
60
|
+
- Pomegranate Engine Architecture: Rind MemTable, Locule immutable containers, and Press compaction.
|
|
61
|
+
- Built-in Quantization: FP32, SQ8, FP16, 1-bit binary quantization, PQ8.
|
|
62
|
+
- Native multi-membrane tenancy and arbitrary payload buffering.
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
setup.py
|
|
4
|
+
pomaidb/__init__.py
|
|
5
|
+
pomaidb/py.typed
|
|
6
|
+
pomaidb/zero_copy.py
|
|
7
|
+
pomaidb.egg-info/PKG-INFO
|
|
8
|
+
pomaidb.egg-info/SOURCES.txt
|
|
9
|
+
pomaidb.egg-info/dependency_links.txt
|
|
10
|
+
pomaidb.egg-info/not-zip-safe
|
|
11
|
+
pomaidb.egg-info/top_level.txt
|
|
12
|
+
pomaidb/lib/libpomai_c.dll
|
|
13
|
+
tests/test_binding.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
pomaidb
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "pomaidb"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Embedded vector database for Edge AI - fast in-process C-API bindings"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = { text = "Apache-2.0" }
|
|
11
|
+
requires-python = ">=3.8"
|
|
12
|
+
authors = [{ name = "PomaiDB Team", email = "info@pomaidb.org" }]
|
|
13
|
+
keywords = ["vector", "embeddings", "hnsw", "database", "embedded", "edge-ai"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 4 - Beta",
|
|
16
|
+
"License :: OSI Approved :: Apache Software License",
|
|
17
|
+
"Programming Language :: Python :: 3",
|
|
18
|
+
"Programming Language :: Python :: 3.8",
|
|
19
|
+
"Programming Language :: Python :: 3.9",
|
|
20
|
+
"Programming Language :: Python :: 3.10",
|
|
21
|
+
"Programming Language :: Python :: 3.11",
|
|
22
|
+
"Programming Language :: Python :: 3.12",
|
|
23
|
+
"Topic :: Database",
|
|
24
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
25
|
+
]
|
|
26
|
+
|
|
27
|
+
[project.urls]
|
|
28
|
+
Homepage = "https://github.com/pomagrenate/pomaidb"
|
|
29
|
+
Repository = "https://github.com/pomagrenate/pomaidb"
|
|
30
|
+
|
|
31
|
+
[tool.setuptools.packages.find]
|
|
32
|
+
where = ["."]
|
|
33
|
+
include = ["pomaidb*"]
|
|
34
|
+
|
|
35
|
+
[tool.setuptools.package-data]
|
|
36
|
+
pomaidb = ["lib/*", "*.dll", "*.so", "*.dylib", "py.typed"]
|
pomaidb-0.1.0/setup.cfg
ADDED
pomaidb-0.1.0/setup.py
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
from setuptools import setup, find_packages
|
|
2
|
+
|
|
3
|
+
setup(
|
|
4
|
+
name="pomaidb",
|
|
5
|
+
version="0.1.0",
|
|
6
|
+
packages=find_packages(),
|
|
7
|
+
package_data={
|
|
8
|
+
"pomaidb": ["lib/*", "*.dll", "*.so", "*.dylib", "py.typed"],
|
|
9
|
+
},
|
|
10
|
+
include_package_data=True,
|
|
11
|
+
zip_safe=False,
|
|
12
|
+
)
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
import shutil
|
|
2
|
+
import tempfile
|
|
3
|
+
import unittest
|
|
4
|
+
import pomaidb
|
|
5
|
+
|
|
6
|
+
class TestPomaiDBBinding(unittest.TestCase):
|
|
7
|
+
def setUp(self):
|
|
8
|
+
self.test_dir = tempfile.mkdtemp(prefix="pomai_py_test_")
|
|
9
|
+
|
|
10
|
+
def tearDown(self):
|
|
11
|
+
shutil.rmtree(self.test_dir, ignore_errors=True)
|
|
12
|
+
|
|
13
|
+
def test_database_class_api(self):
|
|
14
|
+
db = pomaidb.Database.open(self.test_dir, dim=4)
|
|
15
|
+
self.assertIsNotNone(db)
|
|
16
|
+
|
|
17
|
+
# Put vector
|
|
18
|
+
db.put(1, [1.0, 0.0, 0.0, 0.0])
|
|
19
|
+
db.create_membrane("mem1", dim=4)
|
|
20
|
+
db.open_membrane("mem1")
|
|
21
|
+
db.put(2, [0.0, 1.0, 0.0, 0.0], membrane="mem1", payload=b"payload_val", timestamp=42)
|
|
22
|
+
|
|
23
|
+
# Exists
|
|
24
|
+
self.assertTrue(db.exists(1))
|
|
25
|
+
self.assertTrue(db.exists(2, membrane="mem1"))
|
|
26
|
+
|
|
27
|
+
# Get
|
|
28
|
+
v1 = db.get(1)
|
|
29
|
+
self.assertEqual(len(v1), 4)
|
|
30
|
+
self.assertAlmostEqual(v1[0], 1.0, places=3)
|
|
31
|
+
|
|
32
|
+
rec2 = db.get(2, membrane="mem1", with_metadata=True)
|
|
33
|
+
self.assertEqual(rec2.payload, b"payload_val")
|
|
34
|
+
self.assertEqual(rec2.timestamp, 42)
|
|
35
|
+
|
|
36
|
+
# Search
|
|
37
|
+
res = db.search([1.0, 0.0, 0.0, 0.0], topk=2)
|
|
38
|
+
self.assertTrue(len(res) >= 1)
|
|
39
|
+
self.assertEqual(res[0].id, 1)
|
|
40
|
+
|
|
41
|
+
# Stats
|
|
42
|
+
stats = db.get_stats()
|
|
43
|
+
self.assertIn("version", stats)
|
|
44
|
+
|
|
45
|
+
db.flush()
|
|
46
|
+
db.close()
|
|
47
|
+
|
|
48
|
+
def test_functional_c_api(self):
|
|
49
|
+
dir2 = tempfile.mkdtemp(prefix="pomai_py_c_")
|
|
50
|
+
try:
|
|
51
|
+
h = pomaidb.open_db(dir2, dim=4)
|
|
52
|
+
pomaidb.put(h, 10, [0.5, 0.5, 0.5, 0.5])
|
|
53
|
+
self.assertTrue(pomaidb.exists(h, 10))
|
|
54
|
+
v = pomaidb.get(h, 10)
|
|
55
|
+
self.assertEqual(v["id"], 10)
|
|
56
|
+
pomaidb.flush(h)
|
|
57
|
+
pomaidb.close(h)
|
|
58
|
+
finally:
|
|
59
|
+
shutil.rmtree(dir2, ignore_errors=True)
|
|
60
|
+
|
|
61
|
+
if __name__ == "__main__":
|
|
62
|
+
unittest.main()
|