MB-27666: Hierarchy Search

README.md

@@ -24,6 +24,7 @@ A modern indexing + search library in GO
   * [geo spatial search](https://github.com/blevesearch/bleve/blob/master/geo/README.md)
   * approximate k-nearest neighbors via [vector search](https://github.com/blevesearch/bleve/blob/master/docs/vectors.md)
   * [synonym search](https://github.com/blevesearch/bleve/blob/master/docs/synonyms.md)
+  * [hierarchy search](https://github.com/blevesearch/bleve/blob/master/docs/hierarchy.md)
 * [tf-idf](https://github.com/blevesearch/bleve/blob/master/docs/scoring.md#tf-idf) / [bm25](https://github.com/blevesearch/bleve/blob/master/docs/scoring.md#bm25) scoring models
 * Hybrid search: exact + semantic
   * Supports [RRF (Reciprocal Rank Fusion) and RSF (Relative Score Fusion)](docs/score_fusion.md)

document/document.go

@@ -18,6 +18,7 @@
 	"fmt"
 	"reflect"
 
+	"github.com/blevesearch/bleve/v2/search"
 	"github.com/blevesearch/bleve/v2/size"
 	index "github.com/blevesearch/bleve_index_api"
 )
@@ -30,8 +31,9 @@
 }
 
 type Document struct {
-	id               string  `json:"id"`
-	Fields           []Field `json:"fields"`
+	id               string
+	Fields           []Field     `json:"fields"`
+	NestedDocuments  []*Document `json:"nested_documents"`
 	CompositeFields  []*CompositeField
 	StoredFieldsSize uint64
 	indexed          bool
@@ -157,3 +159,34 @@
 func (d *Document) Indexed() bool {
 	return d.indexed
 }
+
+func (d *Document) AddNestedDocument(doc *Document) {
+	d.NestedDocuments = append(d.NestedDocuments, doc)
+}
+
+func (d *Document) NestedFields() search.FieldSet {
+	if len(d.NestedDocuments) == 0 {
+		return nil
+	}
+	fieldSet := search.NewFieldSet()
+	var collectFields func(index.Document)
+	collectFields = func(doc index.Document) {
+		// Add all field names from this nested document
+		doc.VisitFields(func(field index.Field) {
+			fieldSet.AddField(field.Name())
+		})
+		// Recursively collect from this document's nested documents
+		if nd, ok := doc.(index.NestedDocument); ok {
+			nd.VisitNestedDocuments(collectFields)
+		}
+	}
+	// Start collection from nested documents only (not root document)
+	d.VisitNestedDocuments(collectFields)
+	return fieldSet
+}
+
+func (d *Document) VisitNestedDocuments(visitor func(doc index.Document)) {
+	for _, doc := range d.NestedDocuments {
+		visitor(doc)
+	}
+}

index/scorch/introducer.go

@@ -170,6 +170,11 @@ func (s *Scorch) introduceSegment(next *segmentIntroduction) error {
 			newss.deleted = nil
 		}
 
+		// update the deleted bitmap to include any nested/sub-documents as well
+		// if the segment supports that
+		if ns, ok := newss.segment.(segment.NestedSegment); ok {
+			newss.deleted = ns.AddNestedDocuments(newss.deleted)
+		}
 		// check for live size before copying
 		if newss.LiveSize() > 0 {
 			newSnapshot.segment = append(newSnapshot.segment, newss)

index/scorch/scorch.go

@@ -799,6 +799,12 @@ func analyze(d index.Document, fn customAnalyzerPluginInitFunc) {
 			}
 		}
 	})
+	if nd, ok := d.(index.NestedDocument); ok {
+		nd.VisitNestedDocuments(func(doc index.Document) {
+			doc.AddIDField()
+			analyze(doc, fn)
+		})
+	}
 }
 
 func (s *Scorch) AddEligibleForRemoval(epoch uint64) {

index/scorch/snapshot_index.go

@@ -17,7 +17,6 @@ package scorch
 import (
 	"container/heap"
 	"context"
-	"encoding/binary"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -42,9 +41,8 @@ type asynchSegmentResult struct {
 	dict    segment.TermDictionary
 	dictItr segment.DictionaryIterator
 
-	cardinality int
-	index       int
-	docs        *roaring.Bitmap
+	index int
+	docs  *roaring.Bitmap
 
 	thesItr segment.ThesaurusIterator
 
@@ -59,11 +57,11 @@ func init() {
 	var err error
 	lb1, err = lev.NewLevenshteinAutomatonBuilder(1, true)
 	if err != nil {
-		panic(fmt.Errorf("Levenshtein automaton ed1 builder err: %v", err))
+		panic(fmt.Errorf("levenshtein automaton ed1 builder err: %v", err))
 	}
 	lb2, err = lev.NewLevenshteinAutomatonBuilder(2, true)
 	if err != nil {
-		panic(fmt.Errorf("Levenshtein automaton ed2 builder err: %v", err))
+		panic(fmt.Errorf("levenshtein automaton ed2 builder err: %v", err))
 	}
 }
 
@@ -474,7 +472,7 @@ func (is *IndexSnapshot) GetInternal(key []byte) ([]byte, error) {
 func (is *IndexSnapshot) DocCount() (uint64, error) {
 	var rv uint64
 	for _, segment := range is.segment {
-		rv += segment.Count()
+		rv += segment.CountRoot()
 	}
 	return rv, nil
 }
@@ -501,7 +499,7 @@ func (is *IndexSnapshot) Document(id string) (rv index.Document, err error) {
 		return nil, nil
 	}
 
-	docNum, err := docInternalToNumber(next.ID)
+	docNum, err := next.ID.Value()
 	if err != nil {
 		return nil, err
 	}
@@ -571,7 +569,7 @@ func (is *IndexSnapshot) segmentIndexAndLocalDocNumFromGlobal(docNum uint64) (in
 }
 
 func (is *IndexSnapshot) ExternalID(id index.IndexInternalID) (string, error) {
-	docNum, err := docInternalToNumber(id)
+	docNum, err := id.Value()
 	if err != nil {
 		return "", err
 	}
@@ -589,7 +587,7 @@ func (is *IndexSnapshot) ExternalID(id index.IndexInternalID) (string, error) {
 }
 
 func (is *IndexSnapshot) segmentIndexAndLocalDocNum(id index.IndexInternalID) (int, uint64, error) {
-	docNum, err := docInternalToNumber(id)
+	docNum, err := id.Value()
 	if err != nil {
 		return 0, 0, err
 	}
@@ -776,25 +774,6 @@ func (is *IndexSnapshot) recycleTermFieldReader(tfr *IndexSnapshotTermFieldReade
 	is.m2.Unlock()
 }
 
-func docNumberToBytes(buf []byte, in uint64) []byte {
-	if len(buf) != 8 {
-		if cap(buf) >= 8 {
-			buf = buf[0:8]
-		} else {
-			buf = make([]byte, 8)
-		}
-	}
-	binary.BigEndian.PutUint64(buf, in)
-	return buf
-}
-
-func docInternalToNumber(in index.IndexInternalID) (uint64, error) {
-	if len(in) != 8 {
-		return 0, fmt.Errorf("wrong len for IndexInternalID: %q", in)
-	}
-	return binary.BigEndian.Uint64(in), nil
-}
-
 func (is *IndexSnapshot) documentVisitFieldTermsOnSegment(
 	segmentIndex int, localDocNum uint64, fields []string, cFields []string,
 	visitor index.DocValueVisitor, dvs segment.DocVisitState) (
@@ -897,7 +876,7 @@ func (dvr *DocValueReader) BytesRead() uint64 {
 func (dvr *DocValueReader) VisitDocValues(id index.IndexInternalID,
 	visitor index.DocValueVisitor,
 ) (err error) {
-	docNum, err := docInternalToNumber(id)
+	docNum, err := id.Value()
 	if err != nil {
 		return err
 	}
@@ -1297,3 +1276,23 @@ func (is *IndexSnapshot) TermFrequencies(field string, limit int, descending boo
 
 	return termFreqs[:limit], nil
 }
+
+// Ancestors returns the ancestor IDs for the given document ID. The prealloc
+// slice can be provided to avoid allocations downstream, and MUST be empty.
+func (i *IndexSnapshot) Ancestors(ID index.IndexInternalID, prealloc []index.AncestorID) ([]index.AncestorID, error) {
+	// get segment and local doc num for the ID
+	seg, ldoc, err := i.segmentIndexAndLocalDocNum(ID)
+	if err != nil {
+		return nil, err
+	}
+	// get ancestors from the segment
+	prealloc = i.segment[seg].Ancestors(ldoc, prealloc)
+	// get global offset for the segment (correcting factor for multi-segment indexes)
+	globalOffset := i.offsets[seg]
+	// adjust ancestors to global doc numbers, not local to segment
+	for idx := range prealloc {
+		prealloc[idx] = prealloc[idx].Add(globalOffset)
+	}
+	// return adjusted ancestors
+	return prealloc, nil
+}

index/scorch/snapshot_index_doc.go

@@ -15,7 +15,6 @@
 package scorch
 
 import (
-	"bytes"
 	"reflect"
 
 	"github.com/RoaringBitmap/roaring/v2"
@@ -49,7 +48,7 @@ func (i *IndexSnapshotDocIDReader) Next() (index.IndexInternalID, error) {
 		next := i.iterators[i.segmentOffset].Next()
 		// make segment number into global number by adding offset
 		globalOffset := i.snapshot.offsets[i.segmentOffset]
-		return docNumberToBytes(nil, uint64(next)+globalOffset), nil
+		return index.NewIndexInternalID(nil, uint64(next)+globalOffset), nil
 	}
 	return nil, nil
 }
@@ -63,7 +62,7 @@ func (i *IndexSnapshotDocIDReader) Advance(ID index.IndexInternalID) (index.Inde
 	if next == nil {
 		return nil, nil
 	}
-	for bytes.Compare(next, ID) < 0 {
+	for next.Compare(ID) < 0 {
 		next, err = i.Next()
 		if err != nil {
 			return nil, err

index/scorch/snapshot_index_tfr.go

@@ -15,7 +15,6 @@
 package scorch
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"reflect"
@@ -94,7 +93,7 @@ func (i *IndexSnapshotTermFieldReader) Next(preAlloced *index.TermFieldDoc) (*in
 			// make segment number into global number by adding offset
 			globalOffset := i.snapshot.offsets[i.segmentOffset]
 			nnum := next.Number()
-			rv.ID = docNumberToBytes(rv.ID, nnum+globalOffset)
+			rv.ID = index.NewIndexInternalID(rv.ID, nnum+globalOffset)
 			i.postingToTermFieldDoc(next, rv)
 
 			i.currID = rv.ID
@@ -146,7 +145,7 @@ func (i *IndexSnapshotTermFieldReader) postingToTermFieldDoc(next segment.Postin
 func (i *IndexSnapshotTermFieldReader) Advance(ID index.IndexInternalID, preAlloced *index.TermFieldDoc) (*index.TermFieldDoc, error) {
 	// FIXME do something better
 	// for now, if we need to seek backwards, then restart from the beginning
-	if i.currPosting != nil && bytes.Compare(i.currID, ID) >= 0 {
+	if i.currPosting != nil && i.currID.Compare(ID) >= 0 {
 		// Check if the TFR is a special unadorned composite optimization.
 		// Such a TFR will NOT have a valid `term` or `field` set, making it
 		// impossible for the TFR to replace itself with a new one.
@@ -171,7 +170,7 @@ func (i *IndexSnapshotTermFieldReader) Advance(ID index.IndexInternalID, preAllo
 			}
 		}
 	}
-	num, err := docInternalToNumber(ID)
+	num, err := ID.Value()
 	if err != nil {
 		return nil, fmt.Errorf("error converting to doc number % x - %v", ID, err)
 	}
@@ -196,7 +195,7 @@ func (i *IndexSnapshotTermFieldReader) Advance(ID index.IndexInternalID, preAllo
 	if preAlloced == nil {
 		preAlloced = &index.TermFieldDoc{}
 	}
-	preAlloced.ID = docNumberToBytes(preAlloced.ID, next.Number()+
+	preAlloced.ID = index.NewIndexInternalID(preAlloced.ID, next.Number()+
 		i.snapshot.offsets[segIndex])
 	i.postingToTermFieldDoc(next, preAlloced)
 	i.currID = preAlloced.ID

index/scorch/snapshot_index_vr.go

@@ -18,7 +18,6 @@
 package scorch
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -96,7 +95,7 @@ func (i *IndexSnapshotVectorReader) Next(preAlloced *index.VectorDoc) (
 			// make segment number into global number by adding offset
 			globalOffset := i.snapshot.offsets[i.segmentOffset]
 			nnum := next.Number()
-			rv.ID = docNumberToBytes(rv.ID, nnum+globalOffset)
+			rv.ID = index.NewIndexInternalID(rv.ID, nnum+globalOffset)
 			rv.Score = float64(next.Score())
 
 			i.currID = rv.ID
@@ -113,7 +112,7 @@ func (i *IndexSnapshotVectorReader) Next(preAlloced *index.VectorDoc) (
 func (i *IndexSnapshotVectorReader) Advance(ID index.IndexInternalID,
 	preAlloced *index.VectorDoc) (*index.VectorDoc, error) {
 
-	if i.currPosting != nil && bytes.Compare(i.currID, ID) >= 0 {
+	if i.currPosting != nil && i.currID.Compare(ID) >= 0 {
 		i2, err := i.snapshot.VectorReader(i.ctx, i.vector, i.field, i.k,
 			i.searchParams, i.eligibleSelector)
 		if err != nil {
@@ -124,7 +123,7 @@ func (i *IndexSnapshotVectorReader) Advance(ID index.IndexInternalID,
 		*i = *(i2.(*IndexSnapshotVectorReader))
 	}
 
-	num, err := docInternalToNumber(ID)
+	num, err := ID.Value()
 	if err != nil {
 		return nil, fmt.Errorf("error converting to doc number % x - %v", ID, err)
 	}
@@ -149,7 +148,7 @@ func (i *IndexSnapshotVectorReader) Advance(ID index.IndexInternalID,
 	if preAlloced == nil {
 		preAlloced = &index.VectorDoc{}
 	}
-	preAlloced.ID = docNumberToBytes(preAlloced.ID, next.Number()+
+	preAlloced.ID = index.NewIndexInternalID(preAlloced.ID, next.Number()+
 		i.snapshot.offsets[segIndex])
 	i.currID = preAlloced.ID
 	i.currPosting = next

index/scorch/snapshot_segment.go

@@ -113,6 +113,19 @@ func (s *SegmentSnapshot) Count() uint64 {
 	return rv
 }
 
+// this counts the root documents in the segment this differs from Count() in that
+// Count() counts all live documents including nested children, whereas this method
+// counts only root live documents
+func (s *SegmentSnapshot) CountRoot() uint64 {
+	var rv uint64
+	if nsb, ok := s.segment.(segment.NestedSegment); ok {
+		rv = nsb.CountRoot(s.deleted)
+	} else {
+		rv = s.Count()
+	}
+	return rv
+}
+
 func (s *SegmentSnapshot) DocNumbers(docIDs []string) (*roaring.Bitmap, error) {
 	rv, err := s.segment.DocNumbers(docIDs)
 	if err != nil {
@@ -361,3 +374,11 @@ func (c *cachedMeta) fetchMeta(field string) (rv interface{}) {
 	c.m.RUnlock()
 	return rv
 }
+
+func (s *SegmentSnapshot) Ancestors(docNum uint64, prealloc []index.AncestorID) []index.AncestorID {
+	nsb, ok := s.segment.(segment.NestedSegment)
+	if !ok {
+		return append(prealloc, index.NewAncestorID(docNum))
+	}
+	return nsb.Ancestors(docNum, prealloc)
+}