[EXP] a simple JSON RPC client for sourmash search

src/sourmash/index.py

@@ -1,4 +1,42 @@
-"An Abstract Base Class for collections of signatures."
+"""An Abstract Base Class for collections of signatures, plus implementations.
+
+APIs and functionality
+----------------------
+
+Index classes support three sets of API functionality -
+
+'select(...)', which selects subsets of signatures based on ksize, moltype,
+and other criteria, including picklists.
+
+'find(...)', and the 'search', 'gather', and 'counter_gather' implementations
+built on top of 'find', which search for signatures that match a query.
+
+'signatures()', which yields all signatures in the Index subject to the
+selection criteria.
+
+Classes defined in this file
+----------------------------
+
+Index - abstract base class for all Index objects.
+
+LinearIndex - simple in-memory storage of signatures.
+
+LazyLinearIndex - lazy selection and linear search of signatures.
+
+ZipFileLinearIndex - simple on-disk storage of signatures.
+
+class MultiIndex - in-memory storage and selection of signatures from multiple
+index objects, using manifests.
+
+class LazyLoadedIndex - lazy-loading wrapper for on-disk indices, using
+manifests. Signatures are kept on disk until requested; only manifests are
+retained, and no open file handles or signatures.
+
+class LazyMultiIndex - lazy-loading wrapper for many on-disk indices.
+Signatures are kept on disk until requested.
+
+CounterGather - an ancillary class returned by the 'counter_gather()' method.
+"""
 
 import os
 import sourmash
@@ -17,6 +55,7 @@
 
 class Index(ABC):
     is_database = False
+    manifest = None
 
     @property
     def location(self):
@@ -390,9 +429,20 @@ def select(self, **kwargs):
 class LazyLinearIndex(Index):
     """An Index for lazy linear search of another database.
 
-    The defining feature of this class is that 'find' is inherited
-    from the base Index class, which does a linear search with
-    signatures().
+    One of the main purposes of this class is to _force_ linear 'find'
+    on index objects. So if this class wraps an SBT, for example, the
+    SBT find method will be overriden with the linear 'find' from the
+    base class. There are very few situations where this is an improvement,
+    so use this class wisely!
+
+    A few notes:
+    * selection criteria defined by 'select' are only executed when
+      signatures are actually requested (hence, 'lazy').
+    * this class stores the provided index 'db' in memory. If you need
+      a class that does lazy loading of signatures from disk and does not
+      store signatures in memory, see LazyLoadedIndex.
+    * if you want efficient in-memory manifest-based selection, consider
+      LazyMultiIndex.
     """
 
     def __init__(self, db, selection_dict={}):
@@ -474,7 +524,7 @@ def __init__(self, zf, *, selection_dict=None,
         if self.manifest is not None:
             assert not self.selection_dict, self.selection_dict
         if self.selection_dict:
-            assert manifest is None
+            assert self.manifest is None
 
     def _load_manifest(self):
         "Load a manifest if one exists"
@@ -783,13 +833,17 @@ def signatures(self):
             yield row['signature']
 
     def signatures_with_location(self):
+        for row in self.manifest.rows:
+            yield row['signature'], row['internal_location']
+
+    def _signatures_with_internal(self):
         """Return an iterator of tuples (ss, location)
 
         CTB note: here, 'internal_location' is the source file for the
         index. This is a special feature of this (in memory) class.
         """
         for row in self.manifest.rows:
-            yield row['signature'], row['internal_location']
+            yield row['signature'], "", row['internal_location']
 
     def __len__(self):
         return len(self.manifest)
@@ -884,3 +938,171 @@ def select(self, **kwargs):
         "Run 'select' on the manifest."
         new_manifest = self.manifest.select_to_manifest(**kwargs)
         return MultiIndex(new_manifest)
+
+
+class LazyLoadedIndex(Index):
+    """Given an index location and a manifest, do select only on the manifest
+    until signatures are actually requested, and only then load the index.
+
+    This class is useful when you have an index object that consume
+    memory when it is loaded (e.g. JSON signature files, or LCA
+    databases) and you want to avoid keeping them in memory.  The
+    downside of using this class is that it will load the signatures
+    from disk every time they are needed (e.g. 'find(...)', 'signatures()').
+
+    Can be used with LazyMultiIndex to support many such indices at once.
+
+    """
+    def __init__(self, filename, manifest):
+        "Create an Index with given filename and manifest."
+        self.filename = filename
+        self.manifest = manifest
+
+    @property
+    def location(self):
+        "the 'location' attribute for this index will be the filename."
+        return self.filename
+
+    def signatures(self):
+        "yield all signatures from the manifest."
+        if not len(self):
+            # nothing in manifest? done!
+            return []
+
+        # ok - something in manifest, let's go get those signatures!
+        picklist = self.manifest.to_picklist()
+        idx = sourmash.load_file_as_index(self.location)
+
+        # convert remaining manifest into picklist
+        idx = idx.select(picklist=picklist)
+
+        # extract signatures.
+        for ss in idx.signatures():
+            yield ss
+
+    def __len__(self):
+        "track index size based on the manifest."
+        return len(self.manifest)
+    __bool__ = __len__
+
+    @classmethod
+    def load(cls, location, *, create_manifest=False):
+        "Load manifest from given location, and then unload."
+        idx = sourmash.load_file_as_index(location)
+        manifest = idx.manifest
+
+        # do we need to create the manifest?
+        if manifest is None:
+            if create_manifest:
+                iter = idx._signatures_with_internal()
+                manifest = CollectionManifest.create_manifest(idx,
+                                                       include_signature=False)
+            else:
+                raise ValueError(f"no manifest on index at {location}")
+
+        # NOTE: index is not retained outside this scope, just location.
+
+        return cls(location, manifest)
+
+    def insert(self, *args):
+        raise NotImplementedError
+
+    def save(self, *args):
+        raise NotImplementedError
+
+    def select(self, **kwargs):
+        "Run 'select' on manifest, return new object with new manifest."
+        manifest = self.manifest
+        new_manifest = manifest.select_to_manifest(**kwargs)
+
+        return LazyLoadedIndex(self.filename, new_manifest)
+
+
+class LazyMultiIndex(Index):
+    """
+    Do lazy selection of multiple index objects w/manifests.
+
+    Maintains a manifest per collection, and touches the index objects
+    only when actual signatures are needed. This permits lazy loading
+    when wrapping Index classes that are on disk, e.g. ZipFileLinearIndex
+    and SBTs.
+
+    Differs from MultiIndex in that only the manifests are held in memory,
+    not any of the signatures.
+
+    CTB: could we get the same functionality from MultiIndex + LazyLoadedIndex?
+    """
+    def __init__(self, index_list, manifest_list):
+        assert len(index_list) == len(manifest_list)
+        self.index_list = index_list
+        self.manifest_list = manifest_list
+
+    def signatures(self):
+        for ss, loc in self.signatures_with_location():
+            yield ss
+
+    def signatures_with_location(self):
+        for idx, manifest in zip(self.index_list, self.manifest_list):
+            # convert manifest to picklist:
+            picklist = manifest.to_picklist()
+
+            # select using picklist:
+            idx_new = idx.select(picklist=picklist)
+
+            # yield all remaining signatures:
+            for ss, loc in idx_new.signatures_with_location():
+                yield ss, loc
+
+    def __len__(self):
+        return sum( [len(m) for m in self.manifest_list] )
+
+    def insert(self, *args):
+        raise NotImplementedError
+
+    @classmethod
+    def load(cls, index_list):
+        """Create a LazyMultiIndex from a loaded list of index objects.
+
+        All index objects must have manifests already.
+        """
+
+        manifest_list = []
+        for idx in index_list:
+            if not idx.manifest:
+                raise ValueError(f"no manifest on {repr(idx)}")
+            manifest_list.append(idx.manifest)
+
+        # create obj!
+        return cls(index_list, manifest_list)
+
+    @classmethod
+    def load_from_pathlist(cls, filename):
+        """Create a LazyMultiIndex from all files listed in a text file.
+
+        Note, this will not currently work for indices without manifests.
+        """
+        from .sourmash_args import (load_pathlist_from_file,
+                                    load_file_as_index)
+        idx_list = []
+
+        file_list = load_pathlist_from_file(filename)
+        for fname in file_list:
+            idx = load_file_as_index(fname)
+            manifest = getattr(idx, 'manifest', None)
+            if manifest is None:
+                raise ValueError(f"index at '{fname}' has no manifest")
+            idx_list.append(idx)
+
+        return cls.load(idx_list)
+
+    def save(self, *args):
+        raise NotImplementedError
+
+    def select(self, **kwargs):
+        "Run 'select' on all manifests."
+        new_manifests = []
+        for idx, manifest in zip(self.index_list, self.manifest_list):
+            new_manifest = manifest.select_to_manifest(**kwargs)
+            new_manifests.append(new_manifest)
+
+        return LazyMultiIndex(self.index_list, new_manifests)

src/sourmash/manifest.py

@@ -2,6 +2,7 @@
 Manifests for collections of signatures.
 """
 import csv
+from collections import defaultdict
 
 from sourmash.picklist import SignaturePicklist
 
@@ -183,3 +184,58 @@ def to_picklist(self):
         picklist.pickset = set(self._md5_set)
 
         return picklist
+
+
+class ManifestOfManifests:
+    # @CTB rename to MultiManifest?
+    def __init__(self, locations, manifests):
+        assert len(locations) == len(manifests)
+        self.locations = locations
+        self.manifests = manifests
+
+    def __len__(self):
+        return sum([ len(m) for m in self.manifests ])
+
+    @classmethod
+    def load_from_sqlite(cls, filename):
+        import sqlite3
+        db = sqlite3.connect(filename)
+        cursor = db.cursor()
+        cursor.execute('SELECT DISTINCT index_location, internal_location, md5, md5short, ksize, moltype, num, scaled, n_hashes, with_abundance, name, filename FROM manifest')
+
+        d = defaultdict(list)
+        rowkeys = 'internal_location, md5, md5short, ksize, moltype, num, scaled, n_hashes, with_abundance, name, filename'.split(', ')
+        print(rowkeys)
+        for result in cursor:
+            loc, *rest = result
+            mrow = dict(zip(rowkeys, rest))
+            d[loc].append(mrow)
+
+        locs = []
+        manifests = []
+        for loc, value in d.items():
+            manifest = CollectionManifest(value)
+            locs.append(loc)
+            manifests.append(manifest)
+
+        return cls(locs, manifests)
+
+    def select_to_manifest(self, **kwargs):
+        new_manifests = []
+        for m in self.manifests:
+            m = m.select_to_manifest(**kwargs)
+            new_manifests.append(m)
+        return ManifestOfManifests(self.locations, new_manifests)
+
+    def locations(self):
+        raise NotImplementedError
+
+    def locations_and_manifests(self):
+        for (l, m) in zip(self.locations, self.manifests):
+            yield l, m
+
+    def __contains__(self, ss):
+        for m in self.manifests:
+            if ss in m:
+                return True
+        return False

src/sourmash/search.py

@@ -2,13 +2,13 @@
 Code for searching collections of signatures.
 """
 from collections import namedtuple
-from enum import Enum
+from enum import IntEnum
 import numpy as np
 
 from .signature import SourmashSignature
 
 
-class SearchType(Enum):
+class SearchType(IntEnum):
     JACCARD = 1
     CONTAINMENT = 2
     MAX_CONTAINMENT = 3
@@ -95,6 +95,7 @@ def __init__(self, search_type, threshold=None):
             require_scaled = True
         self.score_fn = score_fn
         self.require_scaled = require_scaled
+        self.search_type = search_type
 
         if threshold is None:
             threshold = 0