PyMkDB 0.1.0__py3-none-any.whl

src/db/query/full_text_index.py

@@ -0,0 +1,168 @@
+"""
+FullTextIndex — inverted index for one text field of one store.
+
+File format: idx_{service}_{field}_words.txt
+Each line:   {stem}:{id1},{id2},{id3}
+
+RAM-threshold mode
+------------------
+If ram_threshold_bytes > 0 and the index file on disk exceeds that size,
+the in-memory map is evicted.  Pending mutations are buffered and flushed
+to disk before each query and at each explicit save() call.
+"""
+
+import os
+from src.db.query.tokenizer import tokenize
+
+
+class FullTextIndex:
+    def __init__(self, store_path: str, service: str, field: str,
+                 ram_threshold_bytes: int = 0):
+        safe_field = field.replace("/", "_").replace("\\", "_")
+        self.path = os.path.join(store_path, f"idx_{service}_{safe_field}_words.txt")
+        self._map: dict = {}          # stem -> set[rid]  (RAM mode)
+        self._in_ram: bool = True
+        self._ram_threshold: int = ram_threshold_bytes
+        # Disk-mode pending buffers (only used when _in_ram is False)
+        self._add_buf: dict = {}      # stem -> set[rid] to add
+        self._remove_buf: dict = {}   # stem -> set[rid] to remove
+        self.load()
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+
+    def load(self) -> None:
+        """Load the index from disk into _map."""
+        if not os.path.exists(self.path):
+            return
+        with open(self.path, "r", encoding="utf-8") as fh:
+            for raw in fh:
+                line = raw.strip()
+                if not line or ":" not in line:
+                    continue
+                stem, _, ids_str = line.partition(":")
+                ids = set(ids_str.split(",")) if ids_str else set()
+                self._map[stem] = ids
+        self._apply_threshold()
+
+    def save(self) -> None:
+        """Write the index to disk."""
+        if self._in_ram:
+            self._write_file(self._map)
+            self._apply_threshold()
+        else:
+            self._flush_pending()
+
+    def _apply_threshold(self) -> None:
+        """Evict map from RAM if the disk file exceeds the configured threshold."""
+        if self._ram_threshold > 0 and os.path.exists(self.path):
+            if os.path.getsize(self.path) > self._ram_threshold:
+                self._map.clear()
+                self._in_ram = False
+
+    def _write_file(self, data_map: dict) -> None:
+        """Atomically write a stem->ids map to disk."""
+        tmp = self.path + ".tmp"
+        with open(tmp, "w", encoding="utf-8") as fh:
+            for stem, ids in data_map.items():
+                if ids:
+                    fh.write(f"{stem}:{','.join(sorted(ids))}\n")
+        os.replace(tmp, self.path)
+
+    def _load_disk_map(self) -> dict:
+        """Read the disk file and return a stem->set map without touching self._map."""
+        disk_map: dict = {}
+        if os.path.exists(self.path):
+            with open(self.path, "r", encoding="utf-8") as fh:
+                for raw in fh:
+                    line = raw.strip()
+                    if not line or ":" not in line:
+                        continue
+                    stem, _, ids_str = line.partition(":")
+                    disk_map[stem] = set(ids_str.split(",")) if ids_str else set()
+        return disk_map
+
+    def _flush_pending(self) -> None:
+        """Merge add/remove buffers into the disk file and clear them."""
+        if not self._add_buf and not self._remove_buf:
+            return
+        disk_map = self._load_disk_map()
+
+        for stem, ids in self._add_buf.items():
+            disk_map.setdefault(stem, set()).update(ids)
+
+        for stem, ids in self._remove_buf.items():
+            existing = disk_map.get(stem)
+            if existing:
+                existing -= ids
+                if not existing:
+                    del disk_map[stem]
+
+        self._write_file(disk_map)
+        self._add_buf.clear()
+        self._remove_buf.clear()
+
+    # ------------------------------------------------------------------
+    # Mutation
+    # ------------------------------------------------------------------
+
+    def add(self, record_id: str, text_value: str) -> None:
+        """Tokenize text_value and add record_id to every stem's set."""
+        if self._in_ram:
+            for stem in tokenize(text_value):
+                self._map.setdefault(stem, set()).add(record_id)
+        else:
+            for stem in tokenize(text_value):
+                self._add_buf.setdefault(stem, set()).add(record_id)
+
+    def remove(self, record_id: str, text_value: str) -> None:
+        """Remove record_id from every stem's set; prune empty stems."""
+        if self._in_ram:
+            for stem in tokenize(text_value):
+                ids = self._map.get(stem)
+                if ids:
+                    ids.discard(record_id)
+                    if not ids:
+                        del self._map[stem]
+        else:
+            for stem in tokenize(text_value):
+                self._remove_buf.setdefault(stem, set()).add(record_id)
+
+    # ------------------------------------------------------------------
+    # Query
+    # ------------------------------------------------------------------
+
+    def search(self, keywords: list, mode: str = "and") -> set:
+        """
+        Search for records matching the given keywords.
+
+        keywords : list of raw keyword strings (will be tokenised)
+        mode     : "and" (intersection) | "or" (union)
+        Returns a set of record IDs.
+        """
+        if not keywords:
+            return set()
+        if self._in_ram:
+            data_map = self._map
+        else:
+            self._flush_pending()
+            data_map = self._load_disk_map()
+        candidate_sets = []
+        for kw in keywords:
+            ids = set()
+            for stem in tokenize(kw):
+                ids |= data_map.get(stem, set())
+            candidate_sets.append(ids)
+        if not candidate_sets:
+            return set()
+        if mode == "or":
+            result = set()
+            for s in candidate_sets:
+                result |= s
+            return result
+        # AND: intersect across keywords
+        result = candidate_sets[0]
+        for s in candidate_sets[1:]:
+            result = result & s
+        return result

src/db/query/numeric_index.py

@@ -0,0 +1,196 @@
+"""
+NumericIndex — sorted numeric index for one numeric field of one store.
+
+File format: idx_{service}_{field}_numeric.txt
+Each line:   {value}:{id1},{id2}
+File is kept in ascending sorted order by value.
+
+RAM-threshold mode
+------------------
+If ram_threshold_bytes > 0 and the index file on disk exceeds that size,
+the in-memory map is evicted after load/save.  All queries then read
+directly from disk; pending mutations are buffered and flushed to disk
+before each query and at each explicit save() call.
+"""
+
+import bisect
+import os
+
+
+class NumericIndex:
+    def __init__(self, store_path: str, service: str, field: str,
+                 ram_threshold_bytes: int = 0):
+        safe_field = field.replace("/", "_").replace("\\", "_")
+        self.path = os.path.join(store_path, f"idx_{service}_{safe_field}_numeric.txt")
+        self._values: list = []              # sorted float values  (RAM mode)
+        self._map:    dict = {}              # float -> set[rid]    (RAM mode)
+        self._in_ram: bool = True
+        self._ram_threshold: int = ram_threshold_bytes
+        # Disk-mode pending buffers (only used when _in_ram is False)
+        self._write_buf: dict = {}           # val -> set[rid] to add
+        self._remove_buf: list = []          # list of (val, rid) to remove
+        self.load()
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+
+    def load(self) -> None:
+        if not os.path.exists(self.path):
+            return
+        with open(self.path, "r", encoding="utf-8") as fh:
+            for raw in fh:
+                line = raw.strip()
+                if not line or ":" not in line:
+                    continue
+                val_str, _, ids_str = line.partition(":")
+                try:
+                    val = float(val_str)
+                except ValueError:
+                    continue
+                ids = set(ids_str.split(",")) if ids_str else set()
+                self._map[val] = ids
+        self._values = sorted(self._map.keys())
+        self._apply_threshold()
+
+    def save(self) -> None:
+        if self._in_ram:
+            self._write_file(self._map, self._values)
+            self._apply_threshold()
+        else:
+            self._flush_pending()
+
+    def _apply_threshold(self) -> None:
+        """Evict map from RAM if the disk file exceeds the configured threshold."""
+        if self._ram_threshold > 0 and os.path.exists(self.path):
+            if os.path.getsize(self.path) > self._ram_threshold:
+                self._map.clear()
+                self._values.clear()
+                self._in_ram = False
+
+    def _write_file(self, data_map: dict, data_values: list) -> None:
+        """Atomically write a map+sorted-values to disk."""
+        tmp = self.path + ".tmp"
+        with open(tmp, "w", encoding="utf-8") as fh:
+            for val in data_values:
+                ids = data_map.get(val)
+                if ids:
+                    fh.write(f"{val}:{','.join(sorted(ids))}\n")
+        os.replace(tmp, self.path)
+
+    def _load_disk_map(self) -> tuple:
+        """Read the disk file and return (map, sorted_values) without touching self._map."""
+        disk_map: dict = {}
+        if os.path.exists(self.path):
+            with open(self.path, "r", encoding="utf-8") as fh:
+                for raw in fh:
+                    line = raw.strip()
+                    if not line or ":" not in line:
+                        continue
+                    val_str, _, ids_str = line.partition(":")
+                    try:
+                        val = float(val_str)
+                    except ValueError:
+                        continue
+                    disk_map[val] = set(ids_str.split(",")) if ids_str else set()
+        return disk_map, sorted(disk_map.keys())
+
+    def _flush_pending(self) -> None:
+        """Merge write/remove buffers into the disk file and clear them."""
+        if not self._write_buf and not self._remove_buf:
+            return
+        disk_map, disk_values = self._load_disk_map()
+
+        for val, ids in self._write_buf.items():
+            if val not in disk_map:
+                bisect.insort(disk_values, val)
+                disk_map[val] = set()
+            disk_map[val].update(ids)
+
+        for val, rid in self._remove_buf:
+            ids = disk_map.get(val)
+            if ids:
+                ids.discard(rid)
+                if not ids:
+                    del disk_map[val]
+                    i = bisect.bisect_left(disk_values, val)
+                    if i < len(disk_values) and disk_values[i] == val:
+                        disk_values.pop(i)
+
+        self._write_file(disk_map, disk_values)
+        self._write_buf.clear()
+        self._remove_buf.clear()
+
+    # ------------------------------------------------------------------
+    # Mutation
+    # ------------------------------------------------------------------
+
+    def add(self, record_id: str, numeric_value) -> None:
+        val = float(numeric_value)
+        if self._in_ram:
+            if val not in self._map:
+                bisect.insort(self._values, val)
+                self._map[val] = set()
+            self._map[val].add(record_id)
+        else:
+            self._write_buf.setdefault(val, set()).add(record_id)
+
+    def remove(self, record_id: str, numeric_value) -> None:
+        val = float(numeric_value)
+        if self._in_ram:
+            ids = self._map.get(val)
+            if not ids:
+                return
+            ids.discard(record_id)
+            if not ids:
+                del self._map[val]
+                i = bisect.bisect_left(self._values, val)
+                if i < len(self._values) and self._values[i] == val:
+                    self._values.pop(i)
+        else:
+            self._remove_buf.append((val, record_id))
+
+    # ------------------------------------------------------------------
+    # Query
+    # ------------------------------------------------------------------
+
+    def exact_query(self, value) -> set:
+        if self._in_ram:
+            return set(self._map.get(float(value), set()))
+        self._flush_pending()
+        disk_map, _ = self._load_disk_map()
+        return set(disk_map.get(float(value), set()))
+
+    def range_query(self, lo=None, hi=None,
+                    lo_inclusive: bool = True, hi_inclusive: bool = True) -> set:
+        """
+        Return all record IDs whose value falls in [lo, hi] (bounds optional).
+        lo/hi may be None for open-ended ranges.
+        """
+        if self._in_ram:
+            return self._range_from(self._map, self._values, lo, hi, lo_inclusive, hi_inclusive)
+        self._flush_pending()
+        disk_map, disk_values = self._load_disk_map()
+        return self._range_from(disk_map, disk_values, lo, hi, lo_inclusive, hi_inclusive)
+
+    @staticmethod
+    def _range_from(data_map: dict, data_values: list,
+                    lo, hi, lo_inclusive: bool, hi_inclusive: bool) -> set:
+        if lo is None:
+            lo_idx = 0
+        else:
+            lo_f = float(lo)
+            lo_idx = (bisect.bisect_left(data_values, lo_f) if lo_inclusive
+                      else bisect.bisect_right(data_values, lo_f))
+
+        if hi is None:
+            hi_idx = len(data_values)
+        else:
+            hi_f = float(hi)
+            hi_idx = (bisect.bisect_right(data_values, hi_f) if hi_inclusive
+                      else bisect.bisect_left(data_values, hi_f))
+
+        result: set = set()
+        for val in data_values[lo_idx:hi_idx]:
+            result |= data_map.get(val, set())
+        return result

src/db/query/query_engine.py

@@ -0,0 +1,308 @@
+"""
+QueryEngine — declarative dict query dispatcher for one store.
+
+Query syntax (from the plan's Query Syntax Specification):
+  {"id": "x"}                       → fast-path primary index lookup
+  {"id": ["x", "y"]}                → multi-ID fetch
+  {"field": "exact string"}         → full-scan exact text match
+  {"field": ["kw1", "kw2"]}         → FullTextIndex AND search
+  {"field": 2010}                   → NumericIndex exact_query
+  {"field": {">": 1999, "<": 2015}} → NumericIndex range_query
+  Compound dict → AND-intersect all candidate sets
+"""
+
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+OPERATORS = {">=", "<=", ">", "<"}
+
+# Friendly aliases accepted from SDK callers
+_OP_ALIASES = {"gte": ">=", "lte": "<=", "gt": ">", "lt": "<"}
+
+
+class QuerySyntaxError(ValueError):
+    pass
+
+
+class QueryEngine:
+    def __init__(self, store):
+        """
+        Parameters
+        ----------
+        store : src.db.objects.store.store
+            The owning store instance (used for read() and index_manager).
+        """
+        self._store = store
+        self._full_text_indexes: dict = {}   # field_name -> FullTextIndex
+        self._numeric_indexes:   dict = {}   # field_name -> NumericIndex
+
+    # ------------------------------------------------------------------
+    # Index lifecycle
+    # ------------------------------------------------------------------
+
+    def build_indexes(self) -> None:
+        """
+        Instantiate indexes for all fields marked queryable in schema_config.
+        Called from store.setup() after storage is ready.
+
+        If an index file does not exist on disk, the store is scanned and the
+        index is built and saved immediately so it survives the next restart.
+        Fields whose index file exceeds ram_config.index_ram_threshold_bytes are
+        kept on disk only — queries read from the file; mutations are buffered
+        and flushed before each query.
+        """
+        from src.db.query.full_text_index import FullTextIndex
+        from src.db.query.numeric_index import NumericIndex
+
+        schema = getattr(self._store.config, "schema_config", None)
+        if schema is None:
+            return
+
+        ram_cfg = getattr(self._store.config, "ram_config", None)
+        threshold = getattr(ram_cfg, "index_ram_threshold_bytes", 0) if ram_cfg else 0
+
+        store_name = self._store.config.name
+        print(f"[{store_name}] build_indexes: checking {len(schema.fields)} schema field(s)...")
+
+        needs_rebuild: list = []
+        for field_name, field_schema in schema.fields.items():
+            q = getattr(field_schema, "queryable", False)
+            if not q:
+                continue
+            if q in ("full-text", True):
+                idx = FullTextIndex(
+                    store_path=self._store.store_path,
+                    service=self._store.config.name,
+                    field=field_name,
+                    ram_threshold_bytes=threshold,
+                )
+                self._full_text_indexes[field_name] = idx
+                if not os.path.exists(idx.path):
+                    needs_rebuild.append(field_name)
+            elif q == "numeric":
+                idx = NumericIndex(
+                    store_path=self._store.store_path,
+                    service=self._store.config.name,
+                    field=field_name,
+                    ram_threshold_bytes=threshold,
+                )
+                self._numeric_indexes[field_name] = idx
+                if not os.path.exists(idx.path):
+                    needs_rebuild.append(field_name)
+
+        if needs_rebuild:
+            store_name = self._store.config.name
+            print(f"[{store_name}] Index files missing for: {needs_rebuild} — rebuilding from stored data...")
+            for field_name in needs_rebuild:
+                try:
+                    self.rebuild_index(field_name)
+                except Exception as exc:
+                    import traceback
+                    print(f"[{store_name}] ERROR: rebuild_index('{field_name}') failed: {exc}")
+                    traceback.print_exc()
+                    logger.warning("build_indexes: rebuild_index('%s') failed: %s", field_name, exc)
+        else:
+            print(f"[{self._store.config.name}] All index files present — loaded from disk.")
+
+    def on_write(self, record_id: str, flat_delta: dict) -> None:
+        """Update all relevant indexes when a record is written."""
+        for field, idx in self._full_text_indexes.items():
+            if field in flat_delta:
+                idx.add(record_id, str(flat_delta[field]))
+        for field, idx in self._numeric_indexes.items():
+            if field in flat_delta:
+                try:
+                    idx.add(record_id, flat_delta[field])
+                except (ValueError, TypeError):
+                    pass
+
+    def on_delete(self, record_id: str, flat_dict: dict) -> None:
+        """Remove a record from all indexes."""
+        for field, idx in self._full_text_indexes.items():
+            if field in flat_dict:
+                idx.remove(record_id, str(flat_dict[field]))
+        for field, idx in self._numeric_indexes.items():
+            if field in flat_dict:
+                try:
+                    idx.remove(record_id, flat_dict[field])
+                except (ValueError, TypeError):
+                    pass
+
+    def save_all(self) -> None:
+        """Persist all in-memory indexes to disk."""
+        for idx in self._full_text_indexes.values():
+            idx.save()
+        for idx in self._numeric_indexes.values():
+            idx.save()
+
+    def rebuild_index(self, field_name: str) -> None:
+        """
+        Backfill the index for a single field by scanning all log segments.
+        Safe to call from a background thread — reads only, no writes to log.
+        """
+        from src.db.query.full_text_index import FullTextIndex
+        from src.db.query.numeric_index import NumericIndex
+        from src.db.storage import serializer as _ser
+
+        store = self._store
+        if store.log_manager is None or store.index_manager is None:
+            raise RuntimeError("Store storage not initialised.")
+
+        # Determine index type from schema_config
+        schema = getattr(store.config, "schema_config", None)
+        if schema is None or field_name not in schema.fields:
+            raise QuerySyntaxError(f"Field '{field_name}' not in schema_config.")
+        field_schema = schema.fields[field_name]
+        q = getattr(field_schema, "queryable", False)
+
+        # Determine RAM threshold from store config
+        ram_cfg = getattr(store.config, "ram_config", None)
+        threshold = getattr(ram_cfg, "index_ram_threshold_bytes", 0) if ram_cfg else 0
+
+        # Build a fresh in-memory index (start empty, ignoring any stale file)
+        if q in ("full-text", True):
+            new_idx = FullTextIndex(store.store_path, store.config.name, field_name,
+                                    ram_threshold_bytes=threshold)
+            new_idx._map = {}
+            new_idx._in_ram = True          # force RAM mode while backfilling
+            self._full_text_indexes[field_name] = new_idx
+        elif q == "numeric":
+            new_idx = NumericIndex(store.store_path, store.config.name, field_name,
+                                   ram_threshold_bytes=threshold)
+            new_idx._map = {}
+            new_idx._values = []
+            new_idx._in_ram = True          # force RAM mode while backfilling
+            self._numeric_indexes[field_name] = new_idx
+        else:
+            raise QuerySyntaxError(f"Field '{field_name}' is not queryable.")
+
+        # Scan all live records
+        for record_id in store.index_manager.all_record_ids():
+            entry = store.index_manager.get(record_id)
+            if entry is None:
+                continue
+            seg, offset, size = entry
+            try:
+                line = store.log_manager.read(seg, offset, size)
+                _, flat = _ser.deserialize_record(line)
+                if field_name in flat:
+                    new_idx.add(record_id, flat[field_name])
+            except Exception as exc:
+                logger.warning("rebuild_index: skipping %s: %s", record_id, exc)
+
+        new_idx.save()
+        n = len(new_idx._map)
+        print(f"[{store.config.name}] Index for '{field_name}' rebuilt: {n} entries — saved to {new_idx.path}")
+        logger.info("rebuild_index: field '%s' rebuilt (%d entries)", field_name, n)
+
+    # ------------------------------------------------------------------
+    # Query
+    # ------------------------------------------------------------------
+
+    def query(self, filter_dict: dict) -> list:
+        """
+        Execute a declarative query dict and return a sorted list of matching record IDs.
+        Raises QuerySyntaxError for malformed input.
+        """
+        if not isinstance(filter_dict, dict):
+            raise QuerySyntaxError("filter_dict must be a dict")
+
+        # Empty filter → return all record IDs
+        if not filter_dict:
+            idx_mgr = self._store.index_manager
+            if idx_mgr is None:
+                return []
+            return sorted(idx_mgr.all_record_ids())
+
+        # ID fast path
+        if "id" in filter_dict:
+            ids = filter_dict["id"]
+            if isinstance(ids, str):
+                return [ids]
+            if isinstance(ids, list):
+                return list(ids)
+            raise QuerySyntaxError("'id' value must be a string or list of strings")
+
+        candidate_sets: list = []
+
+        for field, value in filter_dict.items():
+            candidates = self._eval_field(field, value)
+            if not candidates:          # short-circuit
+                return []
+            candidate_sets.append(candidates)
+
+        if not candidate_sets:
+            return []
+
+        result = candidate_sets[0]
+        for s in candidate_sets[1:]:
+            result = result & s
+        return sorted(result)
+
+    def _eval_field(self, field: str, value) -> set:
+        """Evaluate one field clause and return a set of matching record IDs."""
+        if isinstance(value, (int, float)):
+            # Exact numeric match
+            idx = self._numeric_indexes.get(field)
+            if idx is None:
+                raise QuerySyntaxError(
+                    f"Field '{field}' has no numeric index. "
+                    "Mark it queryable='numeric' in schema_config."
+                )
+            return idx.exact_query(value)
+
+        if isinstance(value, str):
+            # Exact text match — linear scan
+            result = set()
+            idx_mgr = self._store.index_manager
+            if idx_mgr is None:
+                return result
+            for record_id in idx_mgr.all_record_ids():
+                rec = self._store.read(record_id)
+                if rec and rec.get(field) == value:
+                    result.add(record_id)
+            return result
+
+        if isinstance(value, list):
+            # Partial / tokenised full-text search (AND across keywords)
+            idx = self._full_text_indexes.get(field)
+            if idx is None:
+                raise QuerySyntaxError(
+                    f"Field '{field}' has no full-text index. "
+                    "Mark it queryable='full-text' in schema_config."
+                )
+            return idx.search(value, mode="and")
+
+        if isinstance(value, dict):
+            # Numeric range query — keys are operators, values are numbers
+            idx = self._numeric_indexes.get(field)
+            if idx is None:
+                raise QuerySyntaxError(
+                    f"Field '{field}' has no numeric index. "
+                    "Mark it queryable='numeric' in schema_config."
+                )
+            lo = hi = None
+            lo_inc = hi_inc = True
+            for op, bound in value.items():
+                op = _OP_ALIASES.get(op, op)   # normalise gte/lte/gt/lt
+                if op not in OPERATORS:
+                    raise QuerySyntaxError(
+                        f"Unsupported operator '{op}'. Must be one of {sorted(OPERATORS)}."
+                    )
+                if op == ">=":
+                    lo, lo_inc = bound, True
+                elif op == ">":
+                    lo, lo_inc = bound, False
+                elif op == "<=":
+                    hi, hi_inc = bound, True
+                elif op == "<":
+                    hi, hi_inc = bound, False
+            return idx.range_query(lo=lo, hi=hi, lo_inclusive=lo_inc, hi_inclusive=hi_inc)
+
+        raise QuerySyntaxError(
+            f"Unsupported value type {type(value).__name__!r} for field '{field}'. "
+            "Use int/float (exact numeric), str (exact text), "
+            "list (keyword search), or dict (range query)."
+        )