diff --git a/codna/testrun.py b/codna/testrun.py index 1473bd4..958ac1f 100644 --- a/codna/testrun.py +++ b/codna/testrun.py @@ -162,6 +162,7 @@ def _detect_pytest_suites(repo_dir: str) -> list[str]: continue if _has_pytest_config(sub): suites.append(sub) + continue if name not in _CONTAINER_DIRS: continue try: @@ -344,17 +345,19 @@ def discover_failing_tests(repo_dir: str, test_cmd: str | None = None, *, timeou "budget. Point Codna at the right one with fix.test_command in codna.yaml or " "--test-cmd \"\" (writing JUnit to $CODNA_JUNIT)." ) - # Each suite runs from its own root so its pyproject/pytest.ini applies. Distribute the - # caller's total integer-second budget as evenly as possible without exceeding it. - base_budget, extra_seconds = divmod(timeout_seconds, len(suites)) + # Each suite runs from its own root so its pyproject/pytest.ini applies; the overall budget + # is split evenly (deterministic). The 60s floor is safe: _MAX_SUITES caps the count, so + # floor * suites can never exceed 8*60s = 480s even when timeout_seconds is small. + per_suite = max(60, timeout_seconds // len(suites)) runs = [] for i, suite in enumerate(suites): suite_junit = junit if len(suites) == 1 else f"{os.path.splitext(junit)[0]}-{i}.xml" - argv = [sys.executable, "-m", "pytest", "-q"] - # A container suite must not recursively collect nested roots that run independently. - argv.extend(f"--ignore={candidate}" for candidate in suites if os.path.dirname(candidate) == suite) - argv.append(f"--junitxml={suite_junit}") - runs.append((suite, argv, suite_junit, base_budget + (1 if i < extra_seconds else 0))) + runs.append(( + suite, + [sys.executable, "-m", "pytest", "-q", f"--junitxml={suite_junit}"], + suite_junit, + per_suite, + )) # Actually ENFORCE network denial (bwrap/unshare on Linux CI) — not just record it — so untrusted # test code can't exfiltrate. When fail-closed egress is required (privacy.egress) but no kernel @@ -368,15 +371,14 @@ def discover_failing_tests(repo_dir: str, test_cmd: str | None = None, *, timeou failing: list[str] = [] per_suite_failing: list[tuple[str, list[str]]] = [] - crashed: list[tuple[str, bool]] = [] + crashed: list[str] = [] saw_failure_rc = False last_tail = "" for run_dir, argv, run_junit, budget in runs: - run_env = {**env, "CODNA_JUNIT": run_junit} # pytest exits 1 when tests fail — that is the expected/accepted case here (not an error). - result = Sandbox( - network="deny", timeout_seconds=budget, env=run_env, network_backend=backend - ).run(argv, cwd=run_dir, accepted_exit_codes=(0, 1)) + result = Sandbox(network="deny", timeout_seconds=budget, env=env, network_backend=backend).run( + argv, cwd=run_dir, accepted_exit_codes=(0, 1) + ) if result.timed_out: raise RuntimeError(f"test run timed out after {budget}s") @@ -413,7 +415,7 @@ def discover_failing_tests(repo_dir: str, test_cmd: str | None = None, *, timeou failing.extend(suite_failing) elif result.returncode not in (0,): saw_failure_rc = True - crashed.append((run_dir, os.path.isfile(run_junit))) + crashed.append(run_dir) last_tail = (result.stdout or result.stderr or "")[-1500:] if failing: @@ -425,12 +427,8 @@ def discover_failing_tests(repo_dir: str, test_cmd: str | None = None, *, timeou ] # A suite that crashed without producing JUnit ids is never dropped from the report just # because a sibling suite produced some — its tail is real failure evidence too. - for suite, has_report in crashed: - message = ( - "exited non-zero with JUnit report but no failing test ids" - if has_report else "exited non-zero with no JUnit report" - ) - parts.append(f"{os.path.relpath(suite, repo_dir)}: {message}") + for suite in crashed: + parts.append(f"{os.path.relpath(suite, repo_dir)}: exited non-zero with no JUnit report") involved = len(per_suite_failing) + len(crashed) return ( f"{len(failing)} failing test(s) across {involved} suite(s) — " diff --git a/memengine/native_runtime.py b/memengine/native_runtime.py index 1aa359a..1f91e3c 100644 --- a/memengine/native_runtime.py +++ b/memengine/native_runtime.py @@ -1,3 +1,4 @@ +# ruff: noqa: E702 """NativeRuntime — ctypes shim over the native Telys runtime (Phase 1b). Implements the full `telys.runtime.RuntimeHandle` against `libty_runtime.dylib` (built from @@ -9,6 +10,7 @@ Opaque handles are `c_void_p` the SDK never introspects; `_NativeIndex` owns the handle and frees it on GC. """ + from __future__ import annotations import ctypes @@ -16,11 +18,15 @@ import os import sys import threading +from collections import defaultdict from contextlib import contextmanager from telys.runtime import FORMAT_VERSION as _SDK_FORMAT_VERSION +from telys.runtime import TEXT_COLUMN as _TEXT_COLUMN from telys.runtime import RuntimeHandle from memengine import _validate +from memengine.lexical import LEX_TOKENS_COLUMN as _LEX_TOKENS_COLUMN +from memengine.lexical import LexicalCodeIndex, tokenize_hashes _DTYPE_CODE = {"f32": 0, "fp16": 1, "int8": 2} _CODE_DTYPE = {v: k for k, v in _DTYPE_CODE.items()} @@ -37,7 +43,7 @@ def __init__(self) -> None: self._readers = 0 self._writer = False self._writers_waiting = 0 - self._wthread = None # owning writer thread (write lock is REENTRANT per thread) + self._wthread = None # owning writer thread (write lock is REENTRANT per thread) self._wdepth = 0 @contextmanager @@ -58,7 +64,7 @@ def read(self): def write(self): me = threading.get_ident() with self._cond: - if self._writer and self._wthread == me: # reentrant: same thread already holds write + if self._writer and self._wthread == me: # reentrant: same thread already holds write self._wdepth += 1 else: self._writers_waiting += 1 @@ -81,19 +87,23 @@ def write(self): def _read(fn): """Guard a read-only RuntimeHandle method under the index read lock (parallel readers).""" + @functools.wraps(fn) def w(self, index, *a, **k): with index._rw.read(): return fn(self, index, *a, **k) + return w def _write(fn): """Guard a mutating RuntimeHandle method under the index write lock (exclusive).""" + @functools.wraps(fn) def w(self, index, *a, **k): with index._rw.write(): return fn(self, index, *a, **k) + return w @@ -104,12 +114,34 @@ def _coerce_where(where): if isinstance(where, dict): if len(where) != 1: return None, None - (c, v), = where.items() + ((c, v),) = where.items() return c, v - col = getattr(where, "column", None) # telys.filters.Eq + col = getattr(where, "column", None) # telys.filters.Eq if col is not None: return col, getattr(where, "value", None) return None, None + + +def _native_col_names(index): + """Return the index's cached native schema, computing it once for legacy handles.""" + native_names = getattr(index, "_native_names", None) + if native_names is None: + native_names = [c for c in index.col_names if c != _LEX_TOKENS_COLUMN] + index._native_names = native_names + return native_names + + +def _partition_sort_key(key): + """Deterministic ordering for nullable partition keys of mixed primitive types.""" + if key is None: + return 0, "", "" + if isinstance(key, float) and key != key: + return 1, "float\x00nan", 0 + if isinstance(key, (str, int, float, bytes)): + return 1, type(key).__name__, key + return 1, type(key).__name__, repr(key) + + _LIB = None _LIB_PATH = None @@ -128,6 +160,7 @@ def _installed_lib(stem: str): """Path to a verified `telys runtime install`ed lib (via telys.paths), or None. Stdlib-only; never raises.""" try: from telys.paths import installed_lib_path + return installed_lib_path(stem) except Exception: # noqa: BLE001 — resolution must never break engine load return None @@ -137,7 +170,10 @@ def _under_install_cache(path) -> bool: """True if `path` is inside the verified `telys runtime install` cache ($TELYS_HOME/runtime/...).""" try: from telys.paths import telys_home - return bool(path) and os.path.abspath(path).startswith(os.path.join(telys_home(), "runtime") + os.sep) + + return bool(path) and os.path.abspath(path).startswith( + os.path.join(telys_home(), "runtime") + os.sep + ) except Exception: # noqa: BLE001 return False @@ -146,15 +182,19 @@ def _resolve_native_lib() -> str: p = os.environ.get("TELYS_NATIVE_KERNEL") or os.environ.get("AME_NATIVE_KERNEL") if p: return p - installed = _installed_lib("libty_runtime") # verified `telys runtime install` location (zero config) + installed = _installed_lib( + "libty_runtime" + ) # verified `telys runtime install` location (zero config) if installed: return installed ext = _lib_ext() here = os.path.dirname(__file__) - bundled = os.path.join(here, "_runtime", f"libty_runtime.{ext}") # internal full-engine wheel + bundled = os.path.join(here, "_runtime", f"libty_runtime.{ext}") # internal full-engine wheel if os.path.exists(bundled): return bundled - return os.path.abspath(os.path.join(here, "..", "..", "..", "mojo_build", f"libty_runtime.{ext}")) + return os.path.abspath( + os.path.join(here, "..", "..", "..", "mojo_build", f"libty_runtime.{ext}") + ) def _lib(): @@ -166,77 +206,170 @@ def _lib(): f"Telys native runtime not found at {_LIB_PATH}. Build it: cd mojo_env && pixi run mojo " "build --emit shared-lib ../mojo/engine/ty_runtime.mojo " f"-o ../mojo_build/libty_runtime.{_lib_ext()} " - "(or set TELYS_NATIVE_KERNEL).") + "(or set TELYS_NATIVE_KERNEL)." + ) lib = ctypes.CDLL(_LIB_PATH) lib.ty_runtime_format_version.restype = ctypes.c_int32 lib.ty_index_new.argtypes = [ctypes.c_int32, ctypes.c_int32] lib.ty_index_new.restype = ctypes.c_void_p lib.ty_index_sealed.argtypes = [ctypes.c_void_p] lib.ty_index_sealed.restype = ctypes.c_int32 - lib.ty_index_stats.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_int64), - ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), - ctypes.POINTER(ctypes.c_int32)] + lib.ty_index_stats.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_int64), + ctypes.POINTER(ctypes.c_int32), + ctypes.POINTER(ctypes.c_int32), + ctypes.POINTER(ctypes.c_int32), + ] lib.ty_index_stats.restype = ctypes.c_int32 lib.ty_index_free.argtypes = [ctypes.c_void_p] lib.ty_index_free.restype = None - lib.ty_index_build.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int64, ctypes.c_int32, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_int32] + lib.ty_index_build.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int32, + ] lib.ty_index_build.restype = ctypes.c_int32 - lib.ty_index_search_partition.argtypes = [ctypes.c_void_p, ctypes.c_int64, ctypes.c_void_p, - ctypes.c_int32, ctypes.c_void_p, ctypes.c_void_p] + lib.ty_index_search_partition.argtypes = [ + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.ty_index_search_partition.restype = ctypes.c_int32 - lib.ty_index_search_full.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, - ctypes.c_void_p, ctypes.c_void_p] + lib.ty_index_search_full.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.ty_index_search_full.restype = ctypes.c_int32 - lib.ty_index_scatter_topk.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_void_p, - ctypes.c_int32, ctypes.c_void_p, ctypes.c_void_p] + lib.ty_index_scatter_topk.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.ty_index_scatter_topk.restype = ctypes.c_int32 lib.ty_index_export_base.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] lib.ty_index_export_base.restype = ctypes.c_int32 # ── MVCC delta layer (1b-B) ── - lib.ty_index_snapshot.argtypes = [ctypes.c_void_p]; lib.ty_index_snapshot.restype = ctypes.c_int64 - lib.ty_index_mutated.argtypes = [ctypes.c_void_p]; lib.ty_index_mutated.restype = ctypes.c_int32 - lib.ty_index_delta_count.argtypes = [ctypes.c_void_p]; lib.ty_index_delta_count.restype = ctypes.c_int64 - lib.ty_index_visible_count.argtypes = [ctypes.c_void_p]; lib.ty_index_visible_count.restype = ctypes.c_int64 + lib.ty_index_snapshot.argtypes = [ctypes.c_void_p] + lib.ty_index_snapshot.restype = ctypes.c_int64 + lib.ty_index_mutated.argtypes = [ctypes.c_void_p] + lib.ty_index_mutated.restype = ctypes.c_int32 + lib.ty_index_delta_count.argtypes = [ctypes.c_void_p] + lib.ty_index_delta_count.restype = ctypes.c_int64 + lib.ty_index_visible_count.argtypes = [ctypes.c_void_p] + lib.ty_index_visible_count.restype = ctypes.c_int64 lib.ty_index_visible_ids.argtypes = [ctypes.c_void_p, ctypes.c_void_p] lib.ty_index_visible_ids.restype = ctypes.c_int64 - lib.ty_index_insert.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32] + lib.ty_index_insert.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int32, + ] lib.ty_index_insert.restype = ctypes.c_int64 lib.ty_index_delete.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32] lib.ty_index_delete.restype = ctypes.c_int64 lib.ty_index_export_visible.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] lib.ty_index_export_visible.restype = ctypes.c_int64 - lib.ty_index_search_mvcc.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_void_p, - ctypes.c_int32, ctypes.c_void_p, ctypes.c_void_p] + lib.ty_index_search_mvcc.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.ty_index_search_mvcc.restype = ctypes.c_int32 # ── filter metadata (1b-B B3) ── - lib.ty_index_set_rowmeta.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_int32, ctypes.c_int32] + lib.ty_index_set_rowmeta.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_int32, + ] lib.ty_index_set_rowmeta.restype = ctypes.c_int32 - lib.ty_index_search_filtered.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int64, - ctypes.c_void_p, ctypes.c_int32, ctypes.c_void_p, ctypes.c_void_p] + lib.ty_index_search_filtered.argtypes = [ + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_int32, + ctypes.c_int64, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.ty_index_search_filtered.restype = ctypes.c_int32 - lib.ty_index_match_ids.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int64, - ctypes.c_void_p] + lib.ty_index_match_ids.argtypes = [ + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_int32, + ctypes.c_int64, + ctypes.c_void_p, + ] lib.ty_index_match_ids.restype = ctypes.c_int64 - lib.ty_index_export_rowmeta.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, - ctypes.c_void_p, ctypes.c_void_p] + lib.ty_index_export_rowmeta.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.ty_index_export_rowmeta.restype = ctypes.c_int32 # ── IVF directory (1b-B B4) ── - lib.ty_index_set_ivf.argtypes = [ctypes.c_void_p, ctypes.c_int64, ctypes.c_int32, ctypes.c_int32, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32] + lib.ty_index_set_ivf.argtypes = [ + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_int32, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int32, + ] lib.ty_index_set_ivf.restype = ctypes.c_int32 - lib.ty_index_clear_ivf.argtypes = [ctypes.c_void_p]; lib.ty_index_clear_ivf.restype = ctypes.c_int32 - lib.ty_index_has_ivf.argtypes = [ctypes.c_void_p, ctypes.c_int64]; lib.ty_index_has_ivf.restype = ctypes.c_int32 + lib.ty_index_clear_ivf.argtypes = [ctypes.c_void_p] + lib.ty_index_clear_ivf.restype = ctypes.c_int32 + lib.ty_index_has_ivf.argtypes = [ctypes.c_void_p, ctypes.c_int64] + lib.ty_index_has_ivf.restype = ctypes.c_int32 lib.ty_index_set_ivf_nprobe.argtypes = [ctypes.c_void_p, ctypes.c_int64, ctypes.c_int32] lib.ty_index_set_ivf_nprobe.restype = ctypes.c_int32 - lib.ty_index_search_ivf.argtypes = [ctypes.c_void_p, ctypes.c_int64, ctypes.c_void_p, ctypes.c_int32, - ctypes.c_void_p, ctypes.c_void_p] + lib.ty_index_search_ivf.argtypes = [ + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.ty_index_search_ivf.restype = ctypes.c_int32 # ── tuner decision policy (1b-B B5) ── - lib.ty_tuner_defaults.argtypes = [ctypes.c_void_p, ctypes.c_void_p]; lib.ty_tuner_defaults.restype = ctypes.c_int32 - lib.ty_index_tuner_oversized.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_void_p, ctypes.c_void_p] + lib.ty_tuner_defaults.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + lib.ty_tuner_defaults.restype = ctypes.c_int32 + lib.ty_index_tuner_oversized.argtypes = [ + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.ty_index_tuner_oversized.restype = ctypes.c_int32 fv = int(lib.ty_runtime_format_version()) if fv != _SDK_FORMAT_VERSION: @@ -256,19 +389,31 @@ def __init__(self, handle, dim, dtype): # 1b-B B3: per-row filter metadata (key_id + column value_ids) lives NATIVELY now. The host keeps only # the tiny value<->id VOCAB needed to translate user filter values to/from the interned ids the native # match uses — not the per-row data. keymap == key_val2id (also the base directory key_id scheme). - self.keymap: dict = {} # partition key value -> key_id (== base directory kid; grows for delta keys) - self.key_i2v: dict = {} # key_id -> partition key value (de-intern for materialize) - self.part_rows: dict = {} # partition key value -> base row count (for explain) - self.col_names: list = [] # filter column names, ordered (position == native col index) - self.col_v2i: dict = {} # col -> {value -> value_id} - self.col_i2v: dict = {} # col -> {value_id -> value} (de-intern for columns_for/materialize) + self.keymap: dict = {} # partition key value -> key_id (== base directory kid; grows for delta keys) + self.key_i2v: dict = {} # key_id -> partition key value (de-intern for materialize) + self.part_rows: dict = {} # partition key value -> base row count (for explain) + self.col_names: list = [] # filter column names, ordered (position == native col index) + self._native_names = None + self.col_v2i: dict = {} # col -> {value -> value_id} + self.col_i2v: dict = {} # col -> {value_id -> value} (de-intern for columns_for/materialize) # IVF-within-partition (1b-B B4): the directory (centroids + cluster lists) + probe/gather/rerank are # NATIVE now. faiss kmeans + nprobe calibration stay host/build-time (D-26). The host keeps only a tiny # membership set + (nlist,nprobe) for routing/explain — not the cluster data. self.ivf_keys: set = set() # partition key VALUES that have a native IVF directory - self.ivf_meta: dict = {} # key value -> {"nlist": int, "nprobe": int} (explain only) + self.ivf_meta: dict = {} # key value -> {"nlist": int, "nprobe": int} (explain only) self.exact_crossover = 1 << 62 - self._rw = _RWLock() # parallel readers / exclusive writer over the host vocab/IVF + native state + # Lexical/BM25 state (host-side exact inverted index over retained token hashes). The native engine does + # not yet store token arrays as row metadata, so we keep the logical-id -> token-hash map on the host and + # persist it as a JSON sidecar. This is a documented exact path; a native BM25 index is future work. + self.lexical_tokens: dict[int, list[int]] = {} # logical id -> token hashes + self.plex: dict = {} # partition key -> LexicalCodeIndex + self.plex_ids: set[int] = set() # logical ids indexed in self.plex + self.lexical_dirty: set[int] = set() # ids updated/deleted since last plex build + self.lexical_k1: float = 1.8 + self.lexical_b: float = 1.0 + self._rw = ( + _RWLock() + ) # parallel readers / exclusive writer over the host vocab/IVF + native state # ── interning: user filter values <-> compact integer ids the native match uses ── def _intern_key(self, value): @@ -303,17 +448,23 @@ class _NativeTuner: emits the public `telys.tuning.TuningPlan`. No readable-Python tuner policy ships in the runtime. Expected metrics are informational (measured optionally at apply via faiss, D-26 — never required, never hot-path). The plan is APPLIED by the runtime-agnostic facade (set_exact_crossover + native build_partition_ivf).""" + name = "heuristic" def __init__(self, target_recall=None, ivf_min_rows=None): mr, tm = ctypes.c_int32(), ctypes.c_int32() - _lib().ty_tuner_defaults(ctypes.byref(mr), ctypes.byref(tm)) # policy constants from the kernel + _lib().ty_tuner_defaults( + ctypes.byref(mr), ctypes.byref(tm) + ) # policy constants from the kernel self.ivf_min_rows = int(ivf_min_rows) if ivf_min_rows is not None else int(mr.value) - self.target_recall = float(target_recall) if target_recall is not None else tm.value / 1000.0 + self.target_recall = ( + float(target_recall) if target_recall is not None else tm.value / 1000.0 + ) def tune_collection(self, collection, workload=None, objectives=None): import numpy as np from telys.tuning import TuningPlan, _hash, _k, _workload_digest + idx = collection.idx cons = (objectives or {}).get("constraints", {}) if objectives else {} target = float(cons.get("recall_at_10", cons.get("recall_at_k", self.target_recall))) @@ -321,31 +472,59 @@ def tune_collection(self, collection, workload=None, objectives=None): # idx._h and rebuilds key_i2v under @_write, so two lock-free native calls would race (use-after-free, # and a TOCTOU overflow if `partitions` grows between the stats read and the oversized read). with idx._rw.read(): - n, dim, dt, nparts = ctypes.c_int64(), ctypes.c_int32(), ctypes.c_int32(), ctypes.c_int32() - _lib().ty_index_stats(idx._h, ctypes.byref(n), ctypes.byref(dim), ctypes.byref(dt), ctypes.byref(nparts)) + n, dim, dt, nparts = ( + ctypes.c_int64(), + ctypes.c_int32(), + ctypes.c_int32(), + ctypes.c_int32(), + ) + _lib().ty_index_stats( + idx._h, ctypes.byref(n), ctypes.byref(dim), ctypes.byref(dt), ctypes.byref(nparts) + ) npart = int(nparts.value) crossover = int(getattr(idx, "exact_crossover", 1 << 62)) oversized = [] - if npart: # the IVF-enable DECISION is computed in the kernel + if npart: # the IVF-enable DECISION is computed in the kernel keyids, rows = np.empty(npart, np.int64), np.empty(npart, np.int64) - cnt = int(_lib().ty_index_tuner_oversized(idx._h, int(self.ivf_min_rows), - keyids.ctypes.data_as(ctypes.c_void_p), rows.ctypes.data_as(ctypes.c_void_p))) - oversized = sorted(((idx.key_i2v.get(int(keyids[i])), int(rows[i])) for i in range(cnt)), - key=lambda x: -x[1]) - trace = [f"target recall floor = {target}", - f"IVF for partitions >= {self.ivf_min_rows} rows (KERNEL policy): {len(oversized)} of {npart}", - "per-partition nprobe calibrated to the floor at apply (build-time, faiss; D-26)"] + cnt = int( + _lib().ty_index_tuner_oversized( + idx._h, + int(self.ivf_min_rows), + keyids.ctypes.data_as(ctypes.c_void_p), + rows.ctypes.data_as(ctypes.c_void_p), + ) + ) + oversized = sorted( + ((idx.key_i2v.get(int(keyids[i])), int(rows[i])) for i in range(cnt)), + key=lambda x: -x[1], + ) + trace = [ + f"target recall floor = {target}", + f"IVF for partitions >= {self.ivf_min_rows} rows (KERNEL policy): {len(oversized)} of {npart}", + "per-partition nprobe calibrated to the floor at apply (build-time, faiss; D-26)", + ] if not oversized: trace.append("no oversized partitions — exact contiguous slice scans; no IVF") return TuningPlan( - collection=collection.name, tuner=self.name, target_recall=target, + collection=collection.name, + tuner=self.name, + target_recall=target, exact_crossover_rows=crossover, - ivf={"enabled": bool(oversized), "min_rows": self.ivf_min_rows, "target_recall": target, - "partitions": [{"key": _k(k), "rows": r} for k, r in oversized[:64]]}, - constraints=dict(cons), decision_trace=trace, + ivf={ + "enabled": bool(oversized), + "min_rows": self.ivf_min_rows, + "target_recall": target, + "partitions": [{"key": _k(k), "rows": r} for k, r in oversized[:64]], + }, + constraints=dict(cons), + decision_trace=trace, workload_hash=_hash(_workload_digest(workload)) if workload is not None else "", - expected={"basis": "kernel_policy", "recall_at_10_target": target, - "note": "IVF nprobe is calibrated to this floor at apply; expected metrics are optional"}) + expected={ + "basis": "kernel_policy", + "recall_at_10_target": target, + "note": "IVF nprobe is calibrated to this floor at apply; expected metrics are optional", + }, + ) def choose_plan(self, stats, query): return None @@ -384,12 +563,24 @@ def stats(self, index) -> dict: def _stats(self, index) -> dict: n, dim = ctypes.c_int64(), ctypes.c_int32() dt, parts = ctypes.c_int32(), ctypes.c_int32() - _lib().ty_index_stats(index._h, ctypes.byref(n), ctypes.byref(dim), ctypes.byref(dt), ctypes.byref(parts)) - visible = int(_lib().ty_index_visible_count(index._h)) # base + delta - tombstoned (native MVCC) - return {"n": visible, "dim": int(dim.value), "dtype": _CODE_DTYPE.get(int(dt.value), "f32"), - "partitions": int(parts.value), "build_ms": 0.0, "index_mb": 0.0, - "overhead_vs_raw_vectors": 0.0, "delta_rows": int(_lib().ty_index_delta_count(index._h)), - "embedding_space": None, "runtime": "native"} + _lib().ty_index_stats( + index._h, ctypes.byref(n), ctypes.byref(dim), ctypes.byref(dt), ctypes.byref(parts) + ) + visible = int( + _lib().ty_index_visible_count(index._h) + ) # base + delta - tombstoned (native MVCC) + return { + "n": visible, + "dim": int(dim.value), + "dtype": _CODE_DTYPE.get(int(dt.value), "f32"), + "partitions": int(parts.value), + "build_ms": 0.0, + "index_mb": 0.0, + "overhead_vs_raw_vectors": 0.0, + "delta_rows": int(_lib().ty_index_delta_count(index._h)), + "embedding_space": None, + "runtime": "native", + } def default_tuner(self): # Kernel-native policy (B5): the decision/thresholds live in libty_runtime — the shipped runtime needs @@ -399,7 +590,8 @@ def default_tuner(self): # ── ingest: seal the base into native state (PR-2) ───────────────────────────────────────────── @_write def build(self, index, vecs, iids, keys, key_name, columns, texts=None): - import numpy as np # texts ignored: lexical index is local-runtime only + import numpy as np + _validate.require_handle(index._h) v = _validate.require_finite(np.ascontiguousarray(vecs, np.float32), "base vectors") if v.ndim != 2: @@ -407,101 +599,159 @@ def build(self, index, vecs, iids, keys, key_name, columns, texts=None): n, dim = v.shape _validate.fits_i32(n, "base row count") keys = np.asarray(keys) - order = np.argsort(keys, kind="stable") # group rows by partition key (contiguous layout) + order = np.argsort(keys, kind="stable") # group rows by partition key (contiguous layout) base = np.ascontiguousarray(v[order], np.float32) phys = np.ascontiguousarray(np.asarray(iids, np.int64)[order]) sk = keys[order] - uniq, starts = np.unique(sk, return_index=True) # ascending unique keys + slice starts + uniq, starts = np.unique(sk, return_index=True) # ascending unique keys + slice starts ndir = len(uniq) - dir_key = np.arange(ndir, dtype=np.int64) # native key_id == position in `uniq` + dir_key = np.arange(ndir, dtype=np.int64) # native key_id == position in `uniq` dir_start = np.ascontiguousarray(starts.astype(np.int32)) - dir_len = np.ascontiguousarray(np.append(starts[1:], n).astype(np.int32) - starts.astype(np.int32)) + dir_len = np.ascontiguousarray( + np.append(starts[1:], n).astype(np.int32) - starts.astype(np.int32) + ) rc = _lib().ty_index_build( - index._h, base.ctypes.data_as(ctypes.c_void_p), int(n), int(dim), - phys.ctypes.data_as(ctypes.c_void_p), dir_key.ctypes.data_as(ctypes.c_void_p), - dir_start.ctypes.data_as(ctypes.c_void_p), dir_len.ctypes.data_as(ctypes.c_void_p), int(ndir)) + index._h, + base.ctypes.data_as(ctypes.c_void_p), + int(n), + int(dim), + phys.ctypes.data_as(ctypes.c_void_p), + dir_key.ctypes.data_as(ctypes.c_void_p), + dir_start.ctypes.data_as(ctypes.c_void_p), + dir_len.ctypes.data_as(ctypes.c_void_p), + int(ndir), + ) if rc != 0: raise RuntimeError(f"ty_index_build failed (rc={rc})") index.key_name = key_name # key_id == base directory kid (position in sorted-unique); the vocab grows for delta-only keys later. index.keymap = {(kk.item() if hasattr(kk, "item") else kk): i for i, kk in enumerate(uniq)} index.key_i2v = {i: kv for kv, i in index.keymap.items()} - index.part_rows = {(kk.item() if hasattr(kk, "item") else kk): int(dir_len[i]) for i, kk in enumerate(uniq)} - index.col_names = list((columns or {}).keys()) # column index == position - index.col_v2i = {}; index.col_i2v = {} + index.part_rows = { + (kk.item() if hasattr(kk, "item") else kk): int(dir_len[i]) for i, kk in enumerate(uniq) + } + index.col_names = list((columns or {}).keys()) # column index == position + index._native_names = [c for c in index.col_names if c != _LEX_TOKENS_COLUMN] + index.col_v2i = {} + index.col_i2v = {} # Per-row filter metadata (key_id + column value_ids) is stored NATIVELY (ty_index_set_rowmeta); the # host only interns the value<->id vocab. The actual per-row data never lives in readable Python. ids_list = [int(x) for x in np.asarray(iids, np.int64)] kvals = [(x.item() if hasattr(x, "item") else x) for x in keys] self._set_rowmeta(index, ids_list, kvals, columns) + # Retain token hashes for the host-side lexical index (native engine does not store arrays as rowmeta). + if texts is not None: + for lid, text in zip(ids_list, texts): + index.lexical_tokens[lid] = tokenize_hashes(str(text)) @staticmethod def _set_rowmeta(index, ids, keyvals, columns): """Intern (key value, column values) -> ids and push the per-row metadata into native state.""" import numpy as np + n = len(ids) if n == 0: return lids = np.ascontiguousarray(np.asarray(ids, np.int64)) keyids = np.ascontiguousarray( - np.asarray([index._intern_key(kv.item() if hasattr(kv, "item") else kv) for kv in keyvals], np.int64)) - nc = len(index.col_names) + np.asarray( + [index._intern_key(kv.item() if hasattr(kv, "item") else kv) for kv in keyvals], + np.int64, + ) + ) + # The token-hash array column is retained in self.lexical_tokens, not in scalar rowmeta. + col_names = _native_col_names(index) + nc = len(col_names) colvals = np.full((n, max(nc, 1)), -1, np.int64) cmap = columns or {} - for ci, cname in enumerate(index.col_names): + for ci, cname in enumerate(col_names): if cname in cmap: arr = list(cmap[cname]) for i in range(n): val = arr[i].item() if hasattr(arr[i], "item") else arr[i] + if cname == _TEXT_COLUMN and val is None: + continue colvals[i, ci] = index._intern_col(cname, val) colvals = np.ascontiguousarray(colvals) - _lib().ty_index_set_rowmeta(index._h, lids.ctypes.data_as(ctypes.c_void_p), - keyids.ctypes.data_as(ctypes.c_void_p), - colvals.ctypes.data_as(ctypes.c_void_p), int(nc), int(n)) + _lib().ty_index_set_rowmeta( + index._h, + lids.ctypes.data_as(ctypes.c_void_p), + keyids.ctypes.data_as(ctypes.c_void_p), + colvals.ctypes.data_as(ctypes.c_void_p), + int(nc), + int(n), + ) # ── query: fast non-mutated key/full paths (directory/full scan + IVF) · native FILTERED match+merge # for off-key and all mutated paths. Filter values are interned to ids; the match runs in Mojo. ────── @_read def search(self, index, vector, top_k, where, explain, target_recall): import numpy as np + _validate.require_handle(index._h) - if np.asarray(vector, np.float32).reshape(-1).shape[0] != index.dim: # n*dim contract (anti-OOB) + if ( + np.asarray(vector, np.float32).reshape(-1).shape[0] != index.dim + ): # n*dim contract (anti-OOB) raise ValueError(f"query vector length != index dim {index.dim}") col, val = _coerce_where(where) value = (val.item() if hasattr(val, "item") else val) if col is not None else None + native_names = _native_col_names(index) # Resolve the filter to (mode, col_idx, value_id). Unknown key/value => matches nothing. if col is None: mode, ci, vid = 0, -1, 0 elif col == index.key_name: mode, ci, vid = 1, -1, index.keymap.get(value) - elif col in index.col_names: - mode, ci, vid = 2, index.col_names.index(col), index.col_v2i.get(col, {}).get(value) + elif col in native_names: + mode, ci, vid = ( + 2, + native_names.index(col), + index.col_v2i.get(col, {}).get(value), + ) else: - raise NotImplementedError(f"NativeRuntime: unknown filter column {col!r} (no usable index/column).") + raise NotImplementedError( + f"NativeRuntime: unknown filter column {col!r} (no usable index/column)." + ) # Clamp top-k to [0, visible_count] so the caller-sized output buffer can never be under-allocated and # a huge top_k can't trigger a giant alloc; reject a non-finite query (NaN silently corrupts ranking). k = _validate.clamp_topk(top_k, max(1, int(_lib().ty_index_visible_count(index._h)))) q = _validate.require_finite(np.ascontiguousarray(vector, np.float32), "query vector") mutated = bool(_lib().ty_index_mutated(index._h)) - plan = "FullScanMask" if mode == 0 else "PartitionSliceExactF32" if mode == 1 else "ScatterGatherExact" - if col is not None and vid is None: # unknown key/column value -> no matches + plan = ( + "FullScanMask" + if mode == 0 + else "PartitionSliceExactF32" + if mode == 1 + else "ScatterGatherExact" + ) + if col is not None and vid is None: # unknown key/column value -> no matches ids, sc = np.empty(0, np.int64), np.empty(0, np.float32) - elif (not mutated) and mode == 1 and target_recall < 1.0 and value in index.ivf_keys: # IVF approx - ids, sc = self._ivf_search(index, vid, q, k); plan = "PartitionIVFRerankF32" - elif (not mutated) and mode == 1: # fast exact partition slice (directory) + elif ( + (not mutated) and mode == 1 and target_recall < 1.0 and value in index.ivf_keys + ): # IVF approx + ids, sc = self._ivf_search(index, vid, q, k) + plan = "PartitionIVFRerankF32" + elif (not mutated) and mode == 1: # fast exact partition slice (directory) ids, sc = self._scan(_lib().ty_index_search_partition, index._h, k, q, c_first=int(vid)) - elif (not mutated) and mode == 0: # fast full scan + elif (not mutated) and mode == 0: # fast full scan ids, sc = self._scan(_lib().ty_index_search_full, index._h, k, q) - else: # native filtered match+merge (off-key / mutated) + else: # native filtered match+merge (off-key / mutated) ids, sc = self._filtered(index, mode, ci, vid, q, k) if explain: - ex = {"plan": plan, "exact": plan != "PartitionIVFRerankF32", "fallback": mode != 1, "mvcc": mutated, - "delta_rows": int(_lib().ty_index_delta_count(index._h)), "candidates_reranked": int(len(ids))} + ex = { + "plan": plan, + "exact": plan != "PartitionIVFRerankF32", + "fallback": mode != 1, + "mvcc": mutated, + "delta_rows": int(_lib().ty_index_delta_count(index._h)), + "candidates_reranked": int(len(ids)), + } if mode == 1: - ex["partition_key"] = index.key_name; ex["partition_value"] = value + ex["partition_key"] = index.key_name + ex["partition_value"] = value ex["base_rows"] = index.part_rows.get(value, 0) if plan == "PartitionIVFRerankF32": - ex["nprobe"] = index.ivf_meta[value]["nprobe"]; ex["nlist"] = index.ivf_meta[value]["nlist"] + ex["nprobe"] = index.ivf_meta[value]["nprobe"] + ex["nlist"] = index.ivf_meta[value]["nlist"] elif mode == 2: ex["filter"] = f"{col} == {value!r}" return ids, sc, ex @@ -510,13 +760,20 @@ def search(self, index, vector, top_k, where, explain, target_recall): @staticmethod def _scan(fn, handle, k, q, c_first=None): import numpy as np + out_ids, out_sc = np.empty(max(k, 1), np.int64), np.empty(max(k, 1), np.float32) args = [handle] if c_first is not None: args.append(c_first) - args += [q.ctypes.data_as(ctypes.c_void_p), k, - out_ids.ctypes.data_as(ctypes.c_void_p), out_sc.ctypes.data_as(ctypes.c_void_p)] - cnt = _validate.check_count(fn(*args), out_ids.shape[0]) # tripwire: kernel must not write past the buffer + args += [ + q.ctypes.data_as(ctypes.c_void_p), + k, + out_ids.ctypes.data_as(ctypes.c_void_p), + out_sc.ctypes.data_as(ctypes.c_void_p), + ] + cnt = _validate.check_count( + fn(*args), out_ids.shape[0] + ) # tripwire: kernel must not write past the buffer v = out_ids[:cnt] >= 0 return out_ids[:cnt][v].copy(), out_sc[:cnt][v].copy() @@ -524,68 +781,128 @@ def _scan(fn, handle, k, q, c_first=None): def _filtered(index, mode, ci, vid, q, k): """Native filtered MVCC search: match (key/column id) + skip tombstoned + score delta-or-base + merge.""" import numpy as np + out_ids, out_sc = np.empty(max(k, 1), np.int64), np.empty(max(k, 1), np.float32) cnt = _lib().ty_index_search_filtered( - index._h, int(mode), int(ci), int(vid if vid is not None else -1), - q.ctypes.data_as(ctypes.c_void_p), k, - out_ids.ctypes.data_as(ctypes.c_void_p), out_sc.ctypes.data_as(ctypes.c_void_p)) - cnt = _validate.check_count(cnt, out_ids.shape[0]) # tripwire: bounded by the caller's buffer + index._h, + int(mode), + int(ci), + int(vid if vid is not None else -1), + q.ctypes.data_as(ctypes.c_void_p), + k, + out_ids.ctypes.data_as(ctypes.c_void_p), + out_sc.ctypes.data_as(ctypes.c_void_p), + ) + cnt = _validate.check_count( + cnt, out_ids.shape[0] + ) # tripwire: bounded by the caller's buffer m = out_ids[:cnt] >= 0 return out_ids[:cnt][m].copy(), out_sc[:cnt][m].copy() @_write def insert(self, index, vecs, iids, keys, columns, texts=None): - # Native: store the delta vectors + bump the LSN + clear tombstones. Host: key/column metadata only. - # (texts ignored: the lexical index is local-runtime only.) + # Native: store the delta vectors + bump the LSN + clear tombstones. Host: key/column metadata + lexical + # tokens + dirty tracking for the host-side BM25 index. import numpy as np + _validate.require_handle(index._h) v = _validate.require_finite(np.ascontiguousarray(vecs, np.float32), "insert vectors") - _validate.require_matrix(v, index.dim, "insert vectors") # n*dim contract (anti-OOB) + _validate.require_matrix(v, index.dim, "insert vectors") # n*dim contract (anti-OOB) ids64 = np.ascontiguousarray(np.asarray(iids, np.int64)) ids = [int(x) for x in ids64] _validate.fits_i32(len(ids), "insert count") - _lib().ty_index_insert(index._h, v.ctypes.data_as(ctypes.c_void_p), - ids64.ctypes.data_as(ctypes.c_void_p), int(len(ids))) - self._set_rowmeta(index, ids, list(np.asarray(keys)), columns) # native per-row key_id + col value_ids + _lib().ty_index_insert( + index._h, + v.ctypes.data_as(ctypes.c_void_p), + ids64.ctypes.data_as(ctypes.c_void_p), + int(len(ids)), + ) + self._set_rowmeta( + index, ids, list(np.asarray(keys)), columns + ) # native per-row key_id + col value_ids + if texts is not None: + for lid, text in zip(ids, texts): + index.lexical_tokens[lid] = tokenize_hashes(str(text)) + index.lexical_dirty.add(lid) return ids @_write def update(self, index, iids, vecs, keys, columns, texts=None): - import numpy as np # texts ignored: lexical index is local-runtime only + """Update rows; texts=None explicitly means document text is unchanged. + + Stored raw text and lexical token hashes are retained. Callers must pass texts when content changes. + """ + import numpy as np + _validate.require_handle(index._h) v = _validate.require_finite(np.ascontiguousarray(vecs, np.float32), "update vectors") - _validate.require_matrix(v, index.dim, "update vectors") # n*dim contract (anti-OOB) + _validate.require_matrix(v, index.dim, "update vectors") # n*dim contract (anti-OOB) ids64 = np.ascontiguousarray(np.asarray(iids, np.int64)) ids = [int(x) for x in ids64] _validate.fits_i32(len(ids), "update count") - _lib().ty_index_insert(index._h, v.ctypes.data_as(ctypes.c_void_p), # upsert (latest wins) native - ids64.ctypes.data_as(ctypes.c_void_p), int(len(ids))) - if keys is None: # keep current key (de-intern native key_id) - keyvals = [index.key_i2v.get(kid) for kid in self._row_keyids(index, ids)] + previous_keyids = self._row_keyids(index, ids) if texts is None else None + preserved_texts = None + if texts is None and _TEXT_COLUMN in index.col_names: + preserved_texts = self._texts_for(index, ids) + _lib().ty_index_insert( + index._h, + v.ctypes.data_as(ctypes.c_void_p), # upsert (latest wins) native + ids64.ctypes.data_as(ctypes.c_void_p), + int(len(ids)), + ) + if keys is None: # keep current key (de-intern native key_id) + current_keyids = ( + previous_keyids if previous_keyids is not None else self._row_keyids(index, ids) + ) + keyvals = [index.key_i2v.get(kid) for kid in current_keyids] else: keyvals = list(np.asarray(keys)) + if preserved_texts is not None: + columns = dict(columns or {}) + if _TEXT_COLUMN not in columns: + columns[_TEXT_COLUMN] = preserved_texts self._set_rowmeta(index, ids, keyvals, columns) + if texts is not None: + for lid, text in zip(ids, texts): + index.lexical_tokens[lid] = tokenize_hashes(str(text)) + index.lexical_dirty.add(lid) + else: + for lid, previous_kid, keyval in zip(ids, previous_keyids, keyvals): + key = keyval.item() if hasattr(keyval, "item") else keyval + if index.key_i2v.get(previous_kid) != key: + index.lexical_dirty.add(lid) return ids @staticmethod def _row_keyids(index, ids): import numpy as np + m = len(ids) lids = np.ascontiguousarray(np.asarray(ids, np.int64)) kids = np.empty(max(m, 1), np.int64) - cols = np.empty(max(m, 1) * max(len(index.col_names), 1), np.int64) - _lib().ty_index_export_rowmeta(index._h, lids.ctypes.data_as(ctypes.c_void_p), int(m), - kids.ctypes.data_as(ctypes.c_void_p), cols.ctypes.data_as(ctypes.c_void_p)) + cols = np.empty(max(m, 1) * max(len(_native_col_names(index)), 1), np.int64) + _lib().ty_index_export_rowmeta( + index._h, + lids.ctypes.data_as(ctypes.c_void_p), + int(m), + kids.ctypes.data_as(ctypes.c_void_p), + cols.ctypes.data_as(ctypes.c_void_p), + ) return [int(x) for x in kids[:m]] @_write def delete(self, index, iids): import numpy as np + _validate.require_handle(index._h) ids64 = np.ascontiguousarray(np.asarray([int(x) for x in iids], np.int64)) + ids = [int(x) for x in ids64] _validate.fits_i32(len(ids64), "delete count") _lib().ty_index_delete(index._h, ids64.ctypes.data_as(ctypes.c_void_p), int(len(ids64))) - return int(len(ids64)) # native tombstone hides the row from search/live/visible + for lid in ids: + index.lexical_tokens.pop(lid, None) + index.lexical_dirty.add(lid) + return int(len(ids64)) # native tombstone hides the row from search/live/visible @_read def snapshot(self, index): @@ -596,48 +913,92 @@ def _materialize_visible(self, index): """(vecs[m,dim], ids[m], keys[m], columns{c:[m]}) for the visible set (latest version, tombstoned excluded) — computed NATIVELY (ty_index_export_visible); the host only attaches key/column metadata.""" import numpy as np + n = int(_lib().ty_index_visible_count(index._h)) dim = index.dim vecs = np.empty(max(n, 1) * dim, np.float32) lids = np.empty(max(n, 1), np.int64) if n: - _lib().ty_index_export_visible(index._h, vecs.ctypes.data_as(ctypes.c_void_p), - lids.ctypes.data_as(ctypes.c_void_p)) - V = np.ascontiguousarray(vecs[:n * dim].reshape(n, dim), np.float32) if n else np.empty((0, dim), np.float32) + _lib().ty_index_export_visible( + index._h, vecs.ctypes.data_as(ctypes.c_void_p), lids.ctypes.data_as(ctypes.c_void_p) + ) + V = ( + np.ascontiguousarray(vecs[: n * dim].reshape(n, dim), np.float32) + if n + else np.empty((0, dim), np.float32) + ) ids = [int(x) for x in lids[:n]] # de-intern native key_id + column value_ids -> values via the host vocab - nc = len(index.col_names) + native_names = _native_col_names(index) + nc = len(native_names) keyids = np.empty(max(n, 1), np.int64) colvals = np.empty(max(n, 1) * max(nc, 1), np.int64) - vlids = np.ascontiguousarray(lids[:max(n, 1)], np.int64) # bound contiguous buffer for the C call + vlids = np.ascontiguousarray( + lids[: max(n, 1)], np.int64 + ) # bound contiguous buffer for the C call if n: - _lib().ty_index_export_rowmeta(index._h, vlids.ctypes.data_as(ctypes.c_void_p), - int(n), keyids.ctypes.data_as(ctypes.c_void_p), - colvals.ctypes.data_as(ctypes.c_void_p)) + _lib().ty_index_export_rowmeta( + index._h, + vlids.ctypes.data_as(ctypes.c_void_p), + int(n), + keyids.ctypes.data_as(ctypes.c_void_p), + colvals.ctypes.data_as(ctypes.c_void_p), + ) keys = [index.key_i2v.get(int(keyids[i])) for i in range(n)] - cv = colvals[:n * nc].reshape(n, nc) if (n and nc) else None - cols = {cname: [index.col_i2v.get(cname, {}).get(int(cv[i, ci])) for i in range(n)] - for ci, cname in enumerate(index.col_names)} if (n and nc) else {c: [] for c in index.col_names} + cv = colvals[: n * nc].reshape(n, nc) if (n and nc) else None + cols = ( + { + cname: [index.col_i2v.get(cname, {}).get(int(cv[i, ci])) for i in range(n)] + for ci, cname in enumerate(native_names) + } + if (n and nc) + else {c: [] for c in native_names} + ) + if _LEX_TOKENS_COLUMN in index.col_names: + cols[_LEX_TOKENS_COLUMN] = [index.lexical_tokens.get(lid) for lid in ids] return V, np.asarray(ids, np.int64), keys, cols @_write def compact(self, index): import numpy as np + # Parity with LocalRuntime.compact(): return {rows_before, rows_after, reclaimed, lsn} — the # RuntimeHandle seam contract (bench/test_engine_api.py asserts it). lsn is 0: the native rebuild # seals a FRESH handle whose log starts over (visible_count before/after carries the row accounting). before = int(_lib().ty_index_visible_count(index._h)) if not _lib().ty_index_mutated(index._h): return {"rows_before": before, "rows_after": before, "reclaimed": 0, "lsn": 0} - V, ids, keys, cols = self._materialize_visible(index) # native export (still on the old handle) + V, ids, keys, cols = self._materialize_visible( + index + ) # native export (still on the old handle) + visible_ids = set(int(x) for x in ids) old = index._h - index._h = _lib().ty_index_new(int(index.dim), int(_DTYPE_CODE.get(index.dtype, 0))) # fresh: empty delta/tomb/lsn + index._h = _lib().ty_index_new( + int(index.dim), int(_DTYPE_CODE.get(index.dtype, 0)) + ) # fresh: empty delta/tomb/lsn index.keymap, index.key_i2v, index.part_rows = {}, {}, {} index.col_names, index.col_v2i, index.col_i2v = [], {}, {} - index.ivf_keys, index.ivf_meta = set(), {} # IVF sublayouts invalidated by the re-seal (fresh handle) + index._native_names = None + index.ivf_keys, index.ivf_meta = ( + set(), + {}, + ) # IVF sublayouts invalidated by the re-seal (fresh handle) + # Lexical sublayouts are also invalidated; keep only visible rows' tokens. + index.lexical_tokens = { + lid: toks for lid, toks in index.lexical_tokens.items() if lid in visible_ids + } + index.plex = {} + index.plex_ids = set() + index.lexical_dirty = set() if len(ids): - self.build(index, V, ids, np.asarray(keys, dtype=object), index.key_name, - {c: np.array(v, dtype=object) for c, v in cols.items()}) + self.build( + index, + V, + ids, + np.asarray(keys, dtype=object), + index.key_name, + {c: np.array(v, dtype=object) for c, v in cols.items()}, + ) _lib().ty_index_free(old) after = int(_lib().ty_index_visible_count(index._h)) return {"rows_before": before, "rows_after": after, "reclaimed": before - after, "lsn": 0} @@ -646,84 +1007,295 @@ def compact(self, index): def save(self, index, path): import json import numpy as np + V, ids, keys, cols = self._materialize_visible(index) os.makedirs(path, exist_ok=True) np.save(os.path.join(path, "ty_base.npy"), V) # Coerce numpy scalars (np.int64/np.float64/np.str_) to plain Python so values round-trip exactly # through JSON — otherwise default=str would stringify them and off-key filters miss after reopen. sc = lambda x: x.item() if hasattr(x, "item") else x # noqa: E731 - meta = {"dim": index.dim, "dtype": index.dtype, "key_name": index.key_name, - "ids": [int(x) for x in ids], "keys": [sc(x) for x in keys], - "columns": {c: [sc(x) for x in v] for c, v in cols.items()}} + meta = { + "dim": index.dim, + "dtype": index.dtype, + "key_name": index.key_name, + "ids": [int(x) for x in ids], + "keys": [sc(x) for x in keys], + "columns": {c: [sc(x) for x in v] for c, v in cols.items()}, + } with open(os.path.join(path, "ty_native.json"), "w") as f: json.dump(meta, f) + # Persist retained token hashes so lexical indexes can be rebuilt after reopen. + with open(os.path.join(path, "ty_lexical.json"), "w") as f: + json.dump({str(k): v for k, v in index.lexical_tokens.items()}, f) return path def open_index(self, path): import json import numpy as np + base_path = os.path.join(path, "ty_base.npy") if not os.path.exists(base_path): # fail closed with a clear diagnostic instead of a cryptic FileNotFoundError (D-19): a collection # written by a different runtime (e.g. LocalRuntime's manifest.json/vectors.f32) lands here. raise RuntimeError( f"collection at {path!r} was not written by the native runtime (no ty_base.npy). It may be a " - "LocalRuntime collection — open it with the matching runtime (TELYS_RUNTIME), or rebuild it.") + "LocalRuntime collection — open it with the matching runtime (TELYS_RUNTIME), or rebuild it." + ) V = np.load(base_path, allow_pickle=False) meta = json.load(open(os.path.join(path, "ty_native.json"))) idx = self.new_index(int(meta["dim"]), meta["dtype"]) + lex_path = os.path.join(path, "ty_lexical.json") + if os.path.exists(lex_path): + with open(lex_path) as f: + idx.lexical_tokens = {int(k): v for k, v in json.load(f).items()} ids = np.asarray(meta["ids"], np.int64) if len(ids): - self.build(idx, np.ascontiguousarray(V, np.float32), ids, np.array(meta["keys"], dtype=object), - meta["key_name"], {c: np.array(v, dtype=object) for c, v in meta.get("columns", {}).items()}) + self.build( + idx, + np.ascontiguousarray(V, np.float32), + ids, + np.array(meta["keys"], dtype=object), + meta["key_name"], + {c: np.array(v, dtype=object) for c, v in meta.get("columns", {}).items()}, + ) else: idx.key_name = meta["key_name"] return idx - # lexical (on-device BM25 code index) — not yet ported to the native runtime; use TELYS_RUNTIME=local. - # (No @_write: these raise immediately; keeping the decorator on build_partition_ivf below.) + # lexical (BM25 code index) — host-side exact inverted index over retained token hashes. + # The native engine does not yet store token arrays as row metadata, so we keep the logical-id -> token-hash + # map on the host, persist it as a JSON sidecar, and build/search LexicalCodeIndex instances over it. This is + # a documented exact path; a native BM25 postings index inside libty_runtime is future work. + @_write def build_lexical(self, index, k1=1.8, b=1.0): - raise NotImplementedError("the lexical (BM25) index is not yet supported on the native runtime; " - "use the local runtime (TELYS_RUNTIME=local) for mode='lexical'") + index.plex = {} + index.plex_ids = set() + index.lexical_dirty = set() + index.lexical_k1 = float(k1) + index.lexical_b = float(b) + by_key = self._live_by_key(index) + if not by_key or not index.lexical_tokens: + return {"partitions": 0, "docs": 0} + built = {} + for key in sorted(by_key.keys(), key=_partition_sort_key): + lids = by_key[key] + ids, hashes = [], [] + for lid in lids: + toks = index.lexical_tokens.get(lid) + if toks: + ids.append(lid) + hashes.append(list(toks)) + if ids: + index.plex[key] = LexicalCodeIndex(k1=k1, b=b).build_from_hashes(ids, hashes) + built[key] = len(ids) + index.plex_ids.update(ids) + return {"partitions": len(index.plex), "docs": int(sum(built.values()))} + @_read def search_lexical(self, index, text, top_k, where, explain): - raise NotImplementedError("the lexical (BM25) index is not yet supported on the native runtime; " - "use the local runtime (TELYS_RUNTIME=local) for mode='lexical'") + import numpy as np + + if not index.lexical_tokens: + ids = np.empty(0, np.int64) + sc = np.empty(0, np.float32) + if explain: + return ids, sc, {"plan": "NoLexicalData", "mode": "lexical"} + return ids, sc + + k1 = float(getattr(index, "lexical_k1", 1.8)) + b = float(getattr(index, "lexical_b", 1.0)) + col, val = _coerce_where(where) + value = (val.item() if hasattr(val, "item") else val) if col is not None else None + native_names = _native_col_names(index) + + by_key = self._live_by_key(index) + visible = {lid for lids in by_key.values() for lid in lids} + lid_to_key = {lid: key for key, lids in by_key.items() for lid in lids} + + if col is None: + keys = sorted(set(index.plex.keys()) | set(by_key.keys()), key=_partition_sort_key) + hits = [ + h + for key in keys + for h in self._lex_partition_hits(index, key, text, top_k, k1, b, by_key) + ] + plan = "FullScanLexicalBM25" + expect_key = None + elif col == index.key_name: + hits = self._lex_partition_hits(index, value, text, top_k, k1, b, by_key) + plan = "PartitionLexicalBM25" + expect_key = value + elif col in native_names: + matching = set(self._match_col_value(index, col, value, native_names)) & visible + adhoc = self._lex_adhoc(index, sorted(matching), k1, b) + hits = adhoc.search(text, max(top_k, 1)) if adhoc is not None else [] + plan = "ScatterGatherLexicalBM25" + expect_key = None + else: + raise NotImplementedError(f"NativeRuntime: unknown filter column {col!r}") + + out_i, out_s, seen = [], [], set() + for lid, sc in sorted(hits, key=lambda x: (-x[1], x[0])): + lid = int(lid) + if lid in seen or lid not in visible: + continue + if expect_key is not None and lid_to_key.get(lid) != expect_key: + continue + seen.add(lid) + out_i.append(lid) + out_s.append(float(sc)) + if len(out_i) >= top_k: + break + + ids = np.asarray(out_i, np.int64) + scores = np.asarray(out_s, np.float32) + if explain: + ex = { + "plan": plan, + "mode": "lexical", + "partition_key": index.key_name, + "partition_value": expect_key, + "hits": int(len(ids)), + "mvcc": bool(_lib().ty_index_mutated(index._h)), + } + return ids, scores, ex + return ids, scores + + @staticmethod + def _lex_partition_hits(index, key, text, k, k1, b, by_key): + """BM25 hits for partition `key`: pre-built plex for base rows still current + ad-hoc index for rows + inserted/updated since the plex was built (or never indexed). Mirrors LocalRuntime's base+delta merge. + by_key is the single visibility snapshot shared by the whole lexical query.""" + hits = [] + idx = index.plex.get(key) + if idx is not None: + for lid, sc in idx.search(text, k): + if lid not in index.lexical_dirty and lid in index.plex_ids: + hits.append((int(lid), float(sc))) + live_lids = by_key.get(key, []) + delta_lids = [ + lid for lid in live_lids if lid not in index.plex_ids or lid in index.lexical_dirty + ] + adhoc = NativeRuntime._lex_adhoc(index, delta_lids, k1, b) if delta_lids else None + if adhoc is not None: + hits.extend((int(lid), float(sc)) for lid, sc in adhoc.search(text, max(k, 1))) + return hits + + @staticmethod + def _lex_adhoc(index, lids, k1, b): + """Throwaway LexicalCodeIndex over an arbitrary live logical-id set.""" + ids, hashes = [], [] + for lid in lids: + toks = index.lexical_tokens.get(lid) + if toks: + ids.append(int(lid)) + hashes.append(list(toks)) + if not ids: + return None + return LexicalCodeIndex(k1=k1, b=b).build_from_hashes(ids, hashes) + + @staticmethod + def _live_by_key(index): + """{partition key -> [live logical ids]} at the current snapshot.""" + import numpy as np + + n = int(_lib().ty_index_visible_count(index._h)) + if n == 0: + return {} + lids = np.empty(n, np.int64) + _lib().ty_index_visible_ids(index._h, lids.ctypes.data_as(ctypes.c_void_p)) + native_names = _native_col_names(index) + nc = max(len(native_names), 1) + kids = np.empty(n, np.int64) + cols = np.empty(n * nc, np.int64) + stride = int( + _lib().ty_index_export_rowmeta( + index._h, + lids.ctypes.data_as(ctypes.c_void_p), + int(n), + kids.ctypes.data_as(ctypes.c_void_p), + cols.ctypes.data_as(ctypes.c_void_p), + ) + ) + if stride != len(native_names): + raise RuntimeError(f"rowmeta stride {stride} != host ncols {len(native_names)}") + by_key = defaultdict(list) + for i in range(n): + key = index.key_i2v.get(int(kids[i])) + by_key[key].append(int(lids[i])) + return dict(by_key) + + @staticmethod + def _match_col_value(index, col, value, native_names=None): + """Logical ids whose current visible version has column `col` == `value`.""" + import numpy as np + + native_names = _native_col_names(index) if native_names is None else native_names + if col not in native_names: + return [] + vid = index.col_v2i.get(col, {}).get(value) + if vid is None: + return [] + mode, ci = 2, native_names.index(col) + cap = int(_lib().ty_index_visible_count(index._h)) + if cap == 0: + return [] + out = np.empty(cap, np.int64) + cnt = int( + _lib().ty_index_match_ids( + index._h, int(mode), int(ci), int(vid), out.ctypes.data_as(ctypes.c_void_p) + ) + ) + return [int(x) for x in out[:cnt].tolist()] # compact serve artifacts — the native index handle has no quantized-export path yet; the engine-side # sealer (PartitionedVectorIndex.save(compact=...)) lives in the local runtime. Fail closed (D-19 style). def export_compact(self, index, path, mode): - raise NotImplementedError("compact serve artifacts are not yet supported on the native runtime; " - "export from the local runtime (TELYS_RUNTIME=local)") + raise NotImplementedError( + "compact serve artifacts are not yet supported on the native runtime; " + "export from the local runtime (TELYS_RUNTIME=local)" + ) def open_compact(self, path): - raise NotImplementedError("compact serve artifacts are not yet supported on the native runtime; " - "open with the local runtime (TELYS_RUNTIME=local)") + raise NotImplementedError( + "compact serve artifacts are not yet supported on the native runtime; " + "open with the local runtime (TELYS_RUNTIME=local)" + ) # ── IVF-within-partition (PR-5; faiss clustering host-side, native exact rerank) ─────────────── @_write def build_partition_ivf(self, index, min_rows, target_recall, max_nprobe=64): import faiss import numpy as np - if _lib().ty_index_mutated(index._h): # IVF builds over a clean base; fold deltas first + + if _lib().ty_index_mutated(index._h): # IVF builds over a clean base; fold deltas first self.compact(index) n = int(self._stats(index)["n"]) if n == 0: return {} dim = index.dim - base = np.empty(n * dim, np.float32); phys = np.empty(n, np.int64) - _lib().ty_index_export_base(index._h, base.ctypes.data_as(ctypes.c_void_p), - phys.ctypes.data_as(ctypes.c_void_p)) + base = np.empty(n * dim, np.float32) + phys = np.empty(n, np.int64) + _lib().ty_index_export_base( + index._h, base.ctypes.data_as(ctypes.c_void_p), phys.ctypes.data_as(ctypes.c_void_p) + ) base = base.reshape(n, dim) # de-intern each base row's partition key (native key_id -> value) to group rows per partition - kids = np.empty(n, np.int64); _cv = np.empty(n * max(len(index.col_names), 1), np.int64) - _lib().ty_index_export_rowmeta(index._h, np.ascontiguousarray(phys).ctypes.data_as(ctypes.c_void_p), - int(n), kids.ctypes.data_as(ctypes.c_void_p), - _cv.ctypes.data_as(ctypes.c_void_p)) + kids = np.empty(n, np.int64) + _cv = np.empty(n * max(len(_native_col_names(index)), 1), np.int64) + _lib().ty_index_export_rowmeta( + index._h, + np.ascontiguousarray(phys).ctypes.data_as(ctypes.c_void_p), + int(n), + kids.ctypes.data_as(ctypes.c_void_p), + _cv.ctypes.data_as(ctypes.c_void_p), + ) by_key: dict = {} for p in range(n): by_key.setdefault(index.key_i2v.get(int(kids[p])), []).append(p) - index.ivf_keys, index.ivf_meta = set(), {}; _lib().ty_index_clear_ivf(index._h); built = {} + index.ivf_keys, index.ivf_meta = set(), {} + _lib().ty_index_clear_ivf(index._h) + built = {} for key, prows in by_key.items(): if len(prows) < int(min_rows): continue @@ -734,25 +1306,43 @@ def build_partition_ivf(self, index, min_rows, target_recall, max_nprobe=64): km = faiss.Kmeans(dim, nlist, niter=10, seed=0, verbose=False) km.train(V) cent = np.ascontiguousarray(km.centroids.reshape(nlist, dim), np.float32) - ci = faiss.IndexFlatIP(dim); ci.add(cent) - _, asn = ci.search(V, 1); asn = asn[:, 0] - crows = [[] for _ in range(nlist)] # cluster -> row indices into V - clusters = [[] for _ in range(nlist)] # cluster -> logical ids (for native rerank) + ci = faiss.IndexFlatIP(dim) + ci.add(cent) + _, asn = ci.search(V, 1) + asn = asn[:, 0] + crows = [[] for _ in range(nlist)] # cluster -> row indices into V + clusters = [[] for _ in range(nlist)] # cluster -> logical ids (for native rerank) for i in range(len(prows)): - c = int(asn[i]); crows[c].append(i); clusters[c].append(int(lids[i])) - nprobe = self._calibrate_nprobe(dim, cent, crows, V, lids, float(target_recall), int(max_nprobe)) + c = int(asn[i]) + crows[c].append(i) + clusters[c].append(int(lids[i])) + nprobe = self._calibrate_nprobe( + dim, cent, crows, V, lids, float(target_recall), int(max_nprobe) + ) # store the IVF directory natively: flat cluster lids + nlist+1 offsets + centroids - clus_flat = np.ascontiguousarray(np.asarray([lid for cl in clusters for lid in cl], np.int64)) - coff = np.zeros(nlist + 1, np.int32); acc = 0 + clus_flat = np.ascontiguousarray( + np.asarray([lid for cl in clusters for lid in cl], np.int64) + ) + coff = np.zeros(nlist + 1, np.int32) + acc = 0 for c in range(nlist): - coff[c] = acc; acc += len(clusters[c]) - coff[nlist] = acc; coff = np.ascontiguousarray(coff) + coff[c] = acc + acc += len(clusters[c]) + coff[nlist] = acc + coff = np.ascontiguousarray(coff) cent_c = np.ascontiguousarray(cent, np.float32) - _lib().ty_index_set_ivf(index._h, int(index.keymap[key]), int(nlist), int(nprobe), - cent_c.ctypes.data_as(ctypes.c_void_p), - clus_flat.ctypes.data_as(ctypes.c_void_p), - coff.ctypes.data_as(ctypes.c_void_p), int(len(clus_flat))) - index.ivf_keys.add(key); index.ivf_meta[key] = {"nlist": nlist, "nprobe": nprobe} + _lib().ty_index_set_ivf( + index._h, + int(index.keymap[key]), + int(nlist), + int(nprobe), + cent_c.ctypes.data_as(ctypes.c_void_p), + clus_flat.ctypes.data_as(ctypes.c_void_p), + coff.ctypes.data_as(ctypes.c_void_p), + int(len(clus_flat)), + ) + index.ivf_keys.add(key) + index.ivf_meta[key] = {"nlist": nlist, "nprobe": nprobe} built[key] = {"rows": len(prows), "nlist": nlist, "nprobe": nprobe} return built @@ -763,6 +1353,7 @@ def _calibrate_nprobe(dim, cent, crows, V, lids, target, max_nprobe): wildly over-estimate recall and pick far too small an nprobe (matches LocalRuntime calibration).""" import faiss import numpy as np + nlist = len(cent) rng = np.random.default_rng(0) m = min(128, len(V)) @@ -770,7 +1361,9 @@ def _calibrate_nprobe(dim, cent, crows, V, lids, target, max_nprobe): Q = V[a] + 0.05 * rng.standard_normal((m, dim)).astype(np.float32) Q /= np.linalg.norm(Q, axis=1, keepdims=True) + 1e-12 Q = np.ascontiguousarray(Q, np.float32) - fl = faiss.IndexFlatIP(dim); fl.add(V); _, truth = fl.search(Q, 10) + fl = faiss.IndexFlatIP(dim) + fl.add(V) + _, truth = fl.search(Q, 10) cs = Q @ cent.T for nprobe in [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64]: if nprobe > nlist: @@ -791,16 +1384,22 @@ def _calibrate_nprobe(dim, cent, crows, V, lids, target, max_nprobe): def _ivf_search(self, index, key_id, q, k): """Native IVF query: probe nprobe nearest centroids + gather + exact rerank — all in Mojo.""" import numpy as np + out_ids, out_sc = np.empty(max(k, 1), np.int64), np.empty(max(k, 1), np.float32) cnt = _lib().ty_index_search_ivf( - index._h, int(key_id), q.ctypes.data_as(ctypes.c_void_p), k, - out_ids.ctypes.data_as(ctypes.c_void_p), out_sc.ctypes.data_as(ctypes.c_void_p)) + index._h, + int(key_id), + q.ctypes.data_as(ctypes.c_void_p), + k, + out_ids.ctypes.data_as(ctypes.c_void_p), + out_sc.ctypes.data_as(ctypes.c_void_p), + ) v = out_ids[:cnt] >= 0 return out_ids[:cnt][v].copy(), out_sc[:cnt][v].copy() @_read def partition_ivf_keys(self, index): - return list(index.ivf_keys) # @_read: ivf membership is reset under @_write (compact) + return list(index.ivf_keys) # @_read: ivf membership is reset under @_write (compact) @_write def set_partition_nprobe(self, index, key, nprobe): @@ -813,37 +1412,96 @@ def set_exact_crossover(self, index, rows): index.exact_crossover = int(rows) # ── provenance / enumeration (host-side metadata; MVCC-aware) ────────────────────────────────── + @_read + def texts_for(self, index, ids): + return self._texts_for(index, ids) + + @staticmethod + def _texts_for(index, ids): + import numpy as np + + native_names = _native_col_names(index) + try: + ci = native_names.index(_TEXT_COLUMN) + except ValueError: + return [None for _ in ids] + nc = len(native_names) + if nc == 0: + return [None for _ in ids] + lids = np.ascontiguousarray(np.asarray(ids, np.int64)) + kids = np.empty(len(ids), np.int64) + cv = np.empty(len(ids) * nc, np.int64) + stride = int( + _lib().ty_index_export_rowmeta( + index._h, + lids.ctypes.data_as(ctypes.c_void_p), + int(len(ids)), + kids.ctypes.data_as(ctypes.c_void_p), + cv.ctypes.data_as(ctypes.c_void_p), + ) + ) + if stride != nc: + raise RuntimeError(f"rowmeta stride {stride} != host ncols {nc}") + cv = cv.reshape(len(ids), nc) + v2i = index.col_i2v.get(_TEXT_COLUMN, {}) + out = [] + for i in range(len(ids)): + vid = int(cv[i, ci]) + out.append(v2i.get(vid)) + return out + @_read def columns_for(self, index, ids): import numpy as np + ids = [int(x) for x in ids] - nc = len(index.col_names) + native_names = _native_col_names(index) + nc = len(native_names) if not ids or nc == 0: return [{} for _ in ids] lids = np.ascontiguousarray(np.asarray(ids, np.int64)) kids = np.empty(len(ids), np.int64) cv = np.empty(len(ids) * nc, np.int64) - stride = int(_lib().ty_index_export_rowmeta( - index._h, lids.ctypes.data_as(ctypes.c_void_p), int(len(ids)), - kids.ctypes.data_as(ctypes.c_void_p), cv.ctypes.data_as(ctypes.c_void_p))) - if stride != nc: # invariant: native ncols == host col_names (tripwire) + stride = int( + _lib().ty_index_export_rowmeta( + index._h, + lids.ctypes.data_as(ctypes.c_void_p), + int(len(ids)), + kids.ctypes.data_as(ctypes.c_void_p), + cv.ctypes.data_as(ctypes.c_void_p), + ) + ) + if stride != nc: # invariant: native ncols == host col_names (tripwire) raise RuntimeError(f"rowmeta stride {stride} != host ncols {nc}") cv = cv.reshape(len(ids), nc) - return [{cname: index.col_i2v.get(cname, {}).get(int(cv[i, ci])) for ci, cname in enumerate(index.col_names)} - for i in range(len(ids))] + rows = [] + for i in range(len(ids)): + row = { + cname: index.col_i2v.get(cname, {}).get(int(cv[i, ci])) + for ci, cname in enumerate(native_names) + } + row.pop(_TEXT_COLUMN, None) + rows.append(row) + return rows @_read def live_ids(self, index, where): import numpy as np + col, val = _coerce_where(where) value = (val.item() if hasattr(val, "item") else val) if col is not None else None + native_names = _native_col_names(index) # resolve to (mode, col_idx, value_id); unknown key/value -> no matches if col is None: mode, ci, vid = 0, -1, 0 elif col == index.key_name: mode, ci, vid = 1, -1, index.keymap.get(value) - elif col in index.col_names: - mode, ci, vid = 2, index.col_names.index(col), index.col_v2i.get(col, {}).get(value) + elif col in native_names: + mode, ci, vid = ( + 2, + native_names.index(col), + index.col_v2i.get(col, {}).get(value), + ) else: raise NotImplementedError(f"NativeRuntime: unknown filter column {col!r}") if col is not None and vid is None: @@ -852,17 +1510,25 @@ def live_ids(self, index, where): if cap == 0: return [] out = np.empty(cap, np.int64) - cnt = int(_lib().ty_index_match_ids(index._h, int(mode), int(ci), int(vid), - out.ctypes.data_as(ctypes.c_void_p))) + cnt = int( + _lib().ty_index_match_ids( + index._h, int(mode), int(ci), int(vid), out.ctypes.data_as(ctypes.c_void_p) + ) + ) return [int(x) for x in out[:cnt]] def kernel_info(self) -> dict: try: path = _resolve_native_lib() - src = ("env TELYS_NATIVE_KERNEL" if os.environ.get("TELYS_NATIVE_KERNEL") or os.environ.get("AME_NATIVE_KERNEL") - else "installed" if _under_install_cache(path) - else "bundled" if (os.sep + "_runtime" + os.sep) in (path or "") - else "repo-relative/dev") + src = ( + "env TELYS_NATIVE_KERNEL" + if os.environ.get("TELYS_NATIVE_KERNEL") or os.environ.get("AME_NATIVE_KERNEL") + else "installed" + if _under_install_cache(path) + else "bundled" + if (os.sep + "_runtime" + os.sep) in (path or "") + else "repo-relative/dev" + ) return {"runtime": "native", "source": src, "path": path, "found": os.path.exists(path)} except Exception: # noqa: BLE001 return {"runtime": "native", "source": "native", "path": None, "found": False} diff --git a/memengine/partitioned.py b/memengine/partitioned.py index cf04b85..9a15f25 100644 --- a/memengine/partitioned.py +++ b/memengine/partitioned.py @@ -614,7 +614,8 @@ def delete(self, ids) -> int: def update(self, ids, vectors, keys=None, columns=None, texts=None) -> int: """Update vectors (and optionally the partition key + off-key columns): tombstone-old is implicit via - LSN supersession; the new version is appended to the (possibly new) partition's delta. Returns the LSN.""" + LSN supersession; the new version is appended to the (possibly new) partition's delta. Returns the LSN. + texts=None explicitly means text is unchanged; existing raw text and lexical token hashes are retained.""" self._require_mutable() if texts is not None: columns = self._tokens_column(texts, columns) @@ -824,7 +825,18 @@ def search_lexical(self, text: str, k: int = 10, where=None, explain: bool = Fal delta merge (the wedge); no filter -> union over all partitions; Eq(off-key column) -> an ad-hoc index over the matching LIVE rows. MVCC-correct: base rows superseded by an update are skipped, delta rows carry current tokens, and every hit is re-checked through _visible_key so tombstoned/superseded/ - moved-out rows never surface. Returns (ids, scores) (logical ids) or (ids, scores, explain).""" + moved-out rows never surface. Returns (ids, scores) (logical ids) or (ids, scores, explain). + + Lazy build: if the collection was created with lexical=True and retained token hashes exist, but no + per-partition plex has been fitted yet, search_lexical builds it automatically using (k1, b). Subsequent + mutations are still reflected through the existing base+delta merge, so the plex is not invalidated on + every write.""" + from memengine.lexical import LEX_TOKENS_COLUMN + # Lazy build under the write lock: a lexical-capable collection should just work without requiring an + # explicit build_lexical() call. The race where two readers both trigger a build is harmless because + # build_lexical is idempotent and protected by the same writer-preferring lock used by mutations. + if self.columns.get(LEX_TOKENS_COLUMN) is not None and not self.plex: + self.build_lexical(k1=k1, b=b) with self._rw.read(): if isinstance(where, dict): if len(where) != 1: diff --git a/memengine/runtime.py b/memengine/runtime.py index f971e2e..a57e6c6 100644 --- a/memengine/runtime.py +++ b/memengine/runtime.py @@ -5,16 +5,18 @@ later sub-phase replaces this with a native runtime behind the same interface. `index` arguments are opaque engine handles (today a `PartitionedVectorIndex`). """ + from __future__ import annotations from telys.runtime import FORMAT_VERSION as _SDK_FORMAT_VERSION -from telys.runtime import RuntimeHandle +from telys.runtime import RuntimeHandle, TEXT_COLUMN as _TEXT_COLUMN from memengine.partitioned import PartitionedVectorIndex # Fail closed if the engine's on-disk format diverges from the SDK's public format contract. assert PartitionedVectorIndex.VERSION == _SDK_FORMAT_VERSION, ( - f"engine format v{PartitionedVectorIndex.VERSION} != SDK format v{_SDK_FORMAT_VERSION}") + f"engine format v{PartitionedVectorIndex.VERSION} != SDK format v{_SDK_FORMAT_VERSION}" +) FORMAT_VERSION = _SDK_FORMAT_VERSION @@ -22,9 +24,13 @@ def _under_install_cache(path) -> bool: """True if `path` is inside the verified `telys runtime install` cache ($TELYS_HOME/runtime/...).""" import os + try: from telys.paths import telys_home - return bool(path) and os.path.abspath(path).startswith(os.path.join(telys_home(), "runtime") + os.sep) + + return bool(path) and os.path.abspath(path).startswith( + os.path.join(telys_home(), "runtime") + os.sep + ) except Exception: # noqa: BLE001 return False @@ -35,13 +41,21 @@ class LocalRuntime(RuntimeHandle): FORMAT_VERSION = FORMAT_VERSION # lifecycle - def new_index(self, dim, dtype): return PartitionedVectorIndex(dim, dtype) - def open_index(self, path): return PartitionedVectorIndex.open(path) - def save(self, index, path): return index.save(path) - def sealed(self, index) -> bool: return index._sealed + def new_index(self, dim, dtype): + return PartitionedVectorIndex(dim, dtype) + + def open_index(self, path): + return PartitionedVectorIndex.open(path) + + def save(self, index, path): + return index.save(path) + + def sealed(self, index) -> bool: + return index._sealed # ingest - def attach_embedder(self, index, embedder): return index.attach_embedder(embedder) + def attach_embedder(self, index, embedder): + return index.attach_embedder(embedder) def build(self, index, vecs, iids, keys, key_name, columns, texts=None): return index.build(vecs, iids, keys, key_name=key_name, columns=columns, texts=texts) @@ -54,48 +68,97 @@ def update(self, index, iids, vecs, keys, columns, texts=None): # query def search(self, index, vector, top_k, where, explain, target_recall): - return index.search(vector, top_k, where=where, explain=explain, target_recall=target_recall) - - def columns_for(self, index, ids): return index.columns_for(ids) - def live_ids(self, index, where): return index.live_ids(where) + return index.search( + vector, top_k, where=where, explain=explain, target_recall=target_recall + ) + + def columns_for(self, index, ids): + rows = index.columns_for(ids) + for r in rows: + r.pop(_TEXT_COLUMN, None) + return rows + + def texts_for(self, index, ids): + arr = index.columns.get(_TEXT_COLUMN) + if arr is None: + return [None for _ in ids] + out = [] + for lid in ids: + i = int(lid) + out.append(arr[i] if 0 <= i < len(arr) and arr[i] is not None else None) + return out + + def live_ids(self, index, where): + return index.live_ids(where) # lexical (on-device BM25 code index) - def build_lexical(self, index, k1=1.8, b=1.0): return index.build_lexical(k1=k1, b=b) + def build_lexical(self, index, k1=1.8, b=1.0): + return index.build_lexical(k1=k1, b=b) def search_lexical(self, index, text, top_k, where, explain): return index.search_lexical(text, top_k, where=where, explain=explain) # maintain - def delete(self, index, iids): return index.delete(iids) - def compact(self, index): return index.compact() - def snapshot(self, index): return index.snapshot() - def stats(self, index) -> dict: return index.stats() + def delete(self, index, iids): + return index.delete(iids) + + def compact(self, index): + return index.compact() + + def snapshot(self, index): + return index.snapshot() + + def stats(self, index) -> dict: + return index.stats() # compact serve artifacts — the engine already seals these; LocalRuntime passes straight through - def export_compact(self, index, path, mode): return index.save(path, compact=mode) - def open_compact(self, path): return PartitionedVectorIndex.open(path) + def export_compact(self, index, path, mode): + return index.save(path, compact=mode) + + def open_compact(self, path): + return PartitionedVectorIndex.open(path) # physical tuning state def build_partition_ivf(self, index, min_rows, target_recall, max_nprobe=64): - return index.build_partition_ivf(min_rows=min_rows, target_recall=target_recall, max_nprobe=max_nprobe) + return index.build_partition_ivf( + min_rows=min_rows, target_recall=target_recall, max_nprobe=max_nprobe + ) + + def set_exact_crossover(self, index, rows): + index.exact_crossover = int(rows) - def set_exact_crossover(self, index, rows): index.exact_crossover = int(rows) - def partition_ivf_keys(self, index): return list(index.pivf.keys()) - def set_partition_nprobe(self, index, key, nprobe): index.pivf[key]["nprobe"] = int(nprobe) + def partition_ivf_keys(self, index): + return list(index.pivf.keys()) + + def set_partition_nprobe(self, index, key, nprobe): + index.pivf[key]["nprobe"] = int(nprobe) # default zero-config tuner def default_tuner(self): from memengine.tuning_heuristic import HeuristicTuner + return HeuristicTuner() # diagnostics — kernel resolution, so the SDK/CLI need not import engine internals def kernel_info(self) -> dict: import os from memengine import mojo_backend as mb + path = mb._LIB_PATH - src = ("env TELYS_KERNEL" if os.environ.get("TELYS_KERNEL") - else "env AME_KERNEL" if os.environ.get("AME_KERNEL") - else "installed" if _under_install_cache(path) - else "bundled" if (os.sep + "_runtime" + os.sep) in (path or "") - else "repo-relative/dev") - return {"runtime": "local", "source": src, "path": path, "found": bool(path and os.path.exists(path))} + src = ( + "env TELYS_KERNEL" + if os.environ.get("TELYS_KERNEL") + else "env AME_KERNEL" + if os.environ.get("AME_KERNEL") + else "installed" + if _under_install_cache(path) + else "bundled" + if (os.sep + "_runtime" + os.sep) in (path or "") + else "repo-relative/dev" + ) + return { + "runtime": "local", + "source": src, + "path": path, + "found": bool(path and os.path.exists(path)), + } diff --git a/tests/test_testrun.py b/tests/test_testrun.py index 6c5e204..be3d528 100644 --- a/tests/test_testrun.py +++ b/tests/test_testrun.py @@ -165,77 +165,6 @@ def run(self, argv, *, cwd, accepted_exit_codes=(0,)): assert "bench" in issue and "cli" in issue -def test_detect_pytest_includes_container_and_nested_suites(tmp_path): - (tmp_path / "packages" / "tests").mkdir(parents=True) - (tmp_path / "packages" / "core" / "tests").mkdir(parents=True) - (tmp_path / "packages" / "extra" / "tests").mkdir(parents=True) - assert testrun._detect_pytest_suites(str(tmp_path)) == [ - str(tmp_path / "packages"), - str(tmp_path / "packages" / "core"), - str(tmp_path / "packages" / "extra"), - ] - - -def test_discover_multiple_suites_split_budget_and_set_per_suite_junit(tmp_path, monkeypatch): - (tmp_path / "bench" / "tests").mkdir(parents=True) - (tmp_path / "cli" / "tests").mkdir(parents=True) - monkeypatch.setattr("importlib.util.find_spec", lambda name: object()) - monkeypatch.delenv("CODNA_REQUIRE_EGRESS_DENY", raising=False) - runs = [] - - class _FakeResult: - timed_out = False - returncode = 0 - stdout = "" - stderr = "" - - class _FakeSandbox: - def __init__(self, **kw): - self.kw = kw - - def run(self, argv, *, cwd, accepted_exit_codes=(0,)): - runs.append((self.kw["timeout_seconds"], self.kw["env"]["CODNA_JUNIT"], argv[-1])) - return _FakeResult() - - monkeypatch.setattr(testrun, "Sandbox", _FakeSandbox) - assert testrun.discover_failing_tests(str(tmp_path), timeout_seconds=91) == (None, []) - assert [budget for budget, _, _ in runs] == [46, 45] - assert all(isinstance(budget, int) for budget, _, _ in runs) - assert all(env_path == junit_arg.split("=", 1)[1] for _, env_path, junit_arg in runs) - assert len({env_path for _, env_path, _ in runs}) == 2 - - -def test_discover_container_suite_excludes_independently_run_nested_roots(tmp_path, monkeypatch): - (tmp_path / "packages" / "tests").mkdir(parents=True) - (tmp_path / "packages" / "core" / "tests").mkdir(parents=True) - (tmp_path / "packages" / "extra" / "tests").mkdir(parents=True) - monkeypatch.setattr("importlib.util.find_spec", lambda name: object()) - monkeypatch.delenv("CODNA_REQUIRE_EGRESS_DENY", raising=False) - runs = {} - - class _FakeResult: - timed_out = False - returncode = 0 - stdout = "" - stderr = "" - - class _FakeSandbox: - def __init__(self, **kw): - pass - - def run(self, argv, *, cwd, accepted_exit_codes=(0,)): - runs[cwd] = argv - return _FakeResult() - - monkeypatch.setattr(testrun, "Sandbox", _FakeSandbox) - assert testrun.discover_failing_tests(str(tmp_path)) == (None, []) - container_argv = runs[str(tmp_path / "packages")] - assert f"--ignore={tmp_path / 'packages' / 'core'}" in container_argv - assert f"--ignore={tmp_path / 'packages' / 'extra'}" in container_argv - assert not any(arg.startswith("--ignore=") for arg in runs[str(tmp_path / "packages" / "core")]) - assert not any(arg.startswith("--ignore=") for arg in runs[str(tmp_path / "packages" / "extra")]) - - def test_discover_crashed_suite_is_reported_alongside_junit_failures(tmp_path, monkeypatch): """bench/ crashes with no JUnit while cli/ produces failures: the crash must appear in the issue text even though `failing` is non-empty from the other suite (review: silent drop).""" @@ -264,31 +193,6 @@ def run(self, argv, *, cwd, accepted_exit_codes=(0,)): assert "2 suite(s)" in issue # failing + crashed both counted in the header -def test_discover_crashed_suite_with_empty_junit_does_not_claim_report_is_missing(tmp_path, monkeypatch): - (tmp_path / "bench" / "tests").mkdir(parents=True) - (tmp_path / "cli" / "tests").mkdir(parents=True) - monkeypatch.setattr("importlib.util.find_spec", lambda name: object()) - monkeypatch.delenv("CODNA_REQUIRE_EGRESS_DENY", raising=False) - - class _FakeSandbox: - def __init__(self, **kw): - pass - - def run(self, argv, *, cwd, accepted_exit_codes=(0,)): - junit = argv[-1].split("=", 1)[1] - report = _JUNIT if cwd.endswith("cli") else '' - with open(junit, "w", encoding="utf-8") as fh: - fh.write(report) - return type("R", (), {"timed_out": False, "returncode": 1, "stdout": "", "stderr": ""})() - - monkeypatch.setattr(testrun, "Sandbox", _FakeSandbox) - issue, failing = testrun.discover_failing_tests(str(tmp_path)) - - assert failing == ["pkg.test_mod::test_bad", "pkg.test_mod::test_err"] - assert "bench: exited non-zero with JUnit report but no failing test ids" in issue - assert "no JUnit report" not in issue - - def test_discover_too_many_suites_refuses_with_guidance(tmp_path, monkeypatch): """Past _MAX_SUITES the run cannot stay inside a sane test budget — refuse with the way out.""" for i in range(testrun._MAX_SUITES + 1):