bug fixes and new features.

Author: CascadingRadium

build.go

@@ -172,6 +172,7 @@ func InitSegmentBase(mem []byte, memCRC uint32, chunkMode uint32, numDocs uint64
 		fieldFSTs:           make(map[uint16]*vellum.FST),
 		vecIndexCache:       newVectorIndexCache(),
 		synIndexCache:       newSynonymIndexCache(),
+		nstIndexCache:       newNestedIndexCache(),
 		// following fields gets populated by loadFieldsNew
 		fieldsMap: make(map[string]uint16),
 		dictLocs:  make([]uint64, 0),

dict.go

@@ -104,7 +104,8 @@ func (d *Dictionary) postingsListInit(rv *PostingsList, except *roaring.Bitmap)
 		rv.postings = postings
 	}
 	rv.sb = d.sb
-	rv.except = except
+	// add subdocs to except if needed
+	rv.except = rv.sb.AddSubDocs(except)
 	return rv
 }
 

faiss_vector_posting.go

@@ -547,6 +547,9 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool
 		pos += n
 	}
 
+	// add sub-docs to except if needed
+	except = sb.AddSubDocs(except)
+
 	vecIndex, vecDocIDMap, docVecIDMap, vectorIDsToExclude, err =
 		sb.vecIndexCache.loadOrCreate(fieldIDPlus1, sb.mem[pos:], requiresFiltering,
 			except)

merge.go

@@ -67,6 +67,8 @@ func mergeSegmentBases(segmentBases []*SegmentBase, drops []*roaring.Bitmap, pat
 		_ = os.Remove(path)
 	}
 
+	drops = addSubDocuments(segmentBases, drops)
+
 	// buffer the output
 	br := bufio.NewWriterSize(f, DefaultFileMergerBufferSize)
 
@@ -531,6 +533,54 @@ func mergeStoredAndRemap(segments []*SegmentBase, drops []*roaring.Bitmap,
 		}
 	}
 
+	// calculate new edge list if applicable
+	var newEdgeList map[uint64]uint64
+
+	for segI, segment := range segments {
+		// check for the closure in meantime
+		if isClosed(closeCh) {
+			return 0, nil, seg.ErrClosed
+		}
+		// get the edgeList for this segment
+		edgeList := segment.EdgeList()
+		// if no edgeList, nothing to do
+		if len(edgeList) == 0 {
+			continue
+		}
+		newSegDocNums := rv[segI]
+		for oldChild, oldParent := range edgeList {
+			newParent := newSegDocNums[oldParent]
+			newChild := newSegDocNums[oldChild]
+			if newParent != docDropped &&
+				newChild != docDropped {
+				if newEdgeList == nil {
+					newEdgeList = make(map[uint64]uint64)
+				}
+				newEdgeList[newChild] = newSegDocNums[oldParent]
+			}
+		}
+	}
+
+	// write out the new edge list
+	// first write out the number of entries
+	// which is also the number of valid subDocs
+	// in the merged segment
+	err := binary.Write(w, binary.BigEndian, uint64(len(newEdgeList)))
+	if err != nil {
+		return 0, nil, err
+	}
+	// now the new edge list
+	for child, parent := range newEdgeList {
+		err := binary.Write(w, binary.BigEndian, child)
+		if err != nil {
+			return 0, nil, err
+		}
+		err = binary.Write(w, binary.BigEndian, parent)
+		if err != nil {
+			return 0, nil, err
+		}
+	}
+
 	return storedIndexOffset, rv, nil
 }
 
@@ -614,3 +664,11 @@ func isClosed(closeCh chan struct{}) bool {
 		return false
 	}
 }
+
+func addSubDocuments(sbs []*SegmentBase, drops []*roaring.Bitmap) []*roaring.Bitmap {
+	rv := make([]*roaring.Bitmap, len(drops))
+	for i, sb := range sbs {
+		rv[i] = sb.AddSubDocs(drops[i])
+	}
+	return rv
+}

nested_cache.go

@@ -0,0 +1,149 @@
+package zap
+
+import (
+	"encoding/binary"
+	"sync"
+)
+
+func newNestedIndexCache() *nestedIndexCache {
+	return &nestedIndexCache{}
+}
+
+type nestedIndexCache struct {
+	m sync.RWMutex
+
+	cache *nestedCacheEntry
+}
+
+// Clear clears the synonym cache which would mean that the termID to term map would no longer be available.
+func (nc *nestedIndexCache) Clear() {
+	nc.m.Lock()
+	nc.cache = nil
+	nc.m.Unlock()
+}
+
+// Returns edgeList, ancestry, descendants
+func (nc *nestedIndexCache) loadOrCreate(edgeListOffset uint64, mem []byte) *nestedCacheEntry {
+	nc.m.RLock()
+	if nc.cache != nil {
+		nc.m.RUnlock()
+		return nc.cache
+	}
+	nc.m.RUnlock()
+
+	nc.m.Lock()
+	defer nc.m.Unlock()
+
+	if nc.cache != nil {
+		return nc.cache
+	}
+
+	return nc.createAndCacheLOCKED(edgeListOffset, mem)
+}
+
+// createAndCacheLOCKED creates the synonym index cache for the specified fieldID and caches it.
+func (sc *nestedIndexCache) createAndCacheLOCKED(edgeListOffset uint64, mem []byte) *nestedCacheEntry {
+	// pos stores the current read position
+	pos := edgeListOffset
+	// read number of subDocs which is also the number of edges
+	numEdges := binary.BigEndian.Uint64(mem[pos : pos+8])
+	pos += 8
+	// if no edges or no subDocs, return empty cache
+	if numEdges == 0 {
+		sc.cache = &nestedCacheEntry{}
+		return sc.cache
+	}
+	// edgeList as a map[node]parent
+	edgeList := make(map[uint64]uint64, numEdges)
+	for i := uint64(0); i < numEdges; i++ {
+		child := binary.BigEndian.Uint64(mem[pos : pos+8])
+		pos += 8
+		parent := binary.BigEndian.Uint64(mem[pos : pos+8])
+		pos += 8
+		edgeList[child] = parent
+	}
+	// build ancestry using DFS + memoization
+	ancestry := make(map[uint64][]uint64, numEdges)
+	// memoized DFS
+	var getAncestors func(uint64) []uint64
+	getAncestors = func(node uint64) []uint64 {
+		// if already computed, return
+		if val, ok := ancestry[node]; ok {
+			return val
+		}
+		if parent, ok := edgeList[node]; ok {
+			// compute parent's ancestry + parent itself
+			res := append([]uint64{parent}, getAncestors(parent)...)
+			ancestry[node] = res
+			return res
+		}
+		return nil
+	}
+
+	for child := range edgeList {
+		// only store if non-empty ancestry
+		if v := getAncestors(child); len(v) > 0 {
+			ancestry[child] = v
+		}
+	}
+
+	descendants := make(map[uint64][]uint64, numEdges)
+
+	// Build descendants using ancestry
+	for node := uint64(0); node < numEdges; node++ {
+		if ancestors, ok := ancestry[node]; ok {
+			for _, ancestor := range ancestors {
+				descendants[ancestor] = append(descendants[ancestor], node)
+			}
+		}
+	}
+
+	sc.cache = &nestedCacheEntry{
+		edgeList:    edgeList,
+		ancestry:    ancestry,
+		descendants: descendants,
+	}
+
+	return sc.cache
+}
+
+func (nc *nestedIndexCache) getAncestry(edgeListOffset uint64, mem []byte, docNum uint64) []uint64 {
+	cache := nc.loadOrCreate(edgeListOffset, mem)
+	if cache == nil || cache.ancestry == nil {
+		return nil
+	}
+	return cache.ancestry[docNum]
+}
+
+func (nc *nestedIndexCache) getDescendants(edgeListOffset uint64, mem []byte, docNum uint64) []uint64 {
+	cache := nc.loadOrCreate(edgeListOffset, mem)
+	if cache == nil || cache.ancestry == nil {
+		return nil
+	}
+	return cache.descendants[docNum]
+}
+
+func (nc *nestedIndexCache) getEdgeList(edgeListOffset uint64, mem []byte) map[uint64]uint64 {
+	cache := nc.loadOrCreate(edgeListOffset, mem)
+	if cache == nil || cache.ancestry == nil {
+		return nil
+	}
+	return cache.edgeList
+}
+
+func (nc *nestedIndexCache) getNumSubDocs(edgeListOffset uint64, mem []byte) uint64 {
+	cache := nc.loadOrCreate(edgeListOffset, mem)
+	if cache == nil {
+		return 0
+	}
+	return uint64(len(cache.edgeList))
+}
+
+type nestedCacheEntry struct {
+	// edgeList[node] = parent
+	edgeList map[uint64]uint64
+	// ancestry[node] = list of parents
+	ancestry map[uint64][]uint64
+	// descendants[parent] = list of children
+	descendants map[uint64][]uint64
+}

new.go

@@ -62,7 +62,7 @@ func (*ZapPlugin) newWithChunkMode(results []index.Document,
 		br.Grow(estimateAvgBytesPerDoc * estimateNumResults)
 	}
 
-	s.results = results
+	s.results, s.edgeList = flattenNestedDocuments(results)
 	s.chunkMode = chunkMode
 	s.w = NewCountHashWriter(&br)
 
@@ -72,7 +72,7 @@ func (*ZapPlugin) newWithChunkMode(results []index.Document,
 	}
 
 	sb, err := InitSegmentBase(br.Bytes(), s.w.Sum32(), chunkMode,
-		uint64(len(results)), storedIndexOffset, sectionsIndexOffset)
+		uint64(len(s.results)), storedIndexOffset, sectionsIndexOffset)
 
 	// get the bytes written before the interim's reset() call
 	// write it to the newly formed segment base.
@@ -94,6 +94,9 @@ var interimPool = sync.Pool{New: func() interface{} { return &interim{} }}
 type interim struct {
 	results []index.Document
 
+	// edge list for nested documents: child -> parent
+	edgeList map[int]uint64
+
 	chunkMode uint32
 
 	w *CountHashWriter
@@ -122,6 +125,7 @@ type interim struct {
 
 func (s *interim) reset() (err error) {
 	s.results = nil
+	s.edgeList = nil
 	s.chunkMode = 0
 	s.w = nil
 	for k := range s.FieldsMap {
@@ -408,6 +412,27 @@ func (s *interim) writeStoredFields() (
 		}
 	}
 
+	// write the number of edges in the child -> parent edge list
+	// this will be zero if there are no nested documents
+	// and this number also reflects the number of sub-documents
+	// in the segment
+	err = binary.Write(s.w, binary.BigEndian, uint64(len(s.edgeList)))
+	if err != nil {
+		return 0, err
+	}
+	// write the child -> parent edge list
+	// child and parent are both flattened doc ids
+	for child, parent := range s.edgeList {
+		err = binary.Write(s.w, binary.BigEndian, uint64(child))
+		if err != nil {
+			return 0, err
+		}
+		err = binary.Write(s.w, binary.BigEndian, parent)
+		if err != nil {
+			return 0, err
+		}
+	}
+
 	return storedIndexOffset, nil
 }
 
@@ -437,3 +462,48 @@ func numUvarintBytes(x uint64) (n int) {
 	}
 	return n + 1
 }
+
+// flattenNestedDocuments returns a preorder list of the given documents and all their nested documents,
+// along with a map mapping each flattened index to its parent index (excluding root docs entirely).
+func flattenNestedDocuments(docs []index.Document) ([]index.Document, map[int]uint64) {
+	totalCount := 0
+	for _, doc := range docs {
+		totalCount += countNestedDocuments(doc)
+	}
+
+	flattened := make([]index.Document, 0, totalCount)
+	edgeMap := make(map[int]uint64, totalCount)
+
+	var traverse func(doc index.Document, hasParent bool, parentIdx int)
+	traverse = func(d index.Document, hasParent bool, parentIdx int) {
+		curIdx := len(flattened)
+		flattened = append(flattened, d)
+
+		if hasParent {
+			edgeMap[curIdx] = uint64(parentIdx)
+		}
+
+		if nestedDoc, ok := d.(index.NestedDocument); ok {
+			nestedDoc.VisitNestedDocuments(func(child index.Document) {
+				traverse(child, true, curIdx)
+			})
+		}
+	}
+	// Top-level docs have no parent
+	for _, doc := range docs {
+		traverse(doc, false, 0)
+	}
+	return flattened, edgeMap
+}
+
+// countNestedDocuments returns the total number of docs in preorder,
+// including the parent and all descendants.
+func countNestedDocuments(doc index.Document) int {
+	count := 1 // include this doc
+	if nd, ok := doc.(index.NestedDocument); ok {
+		nd.VisitNestedDocuments(func(child index.Document) {
+			count += countNestedDocuments(child)
+		})
+	}
+	return count
+}

read.go

@@ -41,3 +41,12 @@ func (sb *SegmentBase) getDocStoredOffsets(docNum uint64) (
 
 	return indexOffset, storedOffset, n, metaLen, dataLen
 }
+
+func (sb *SegmentBase) getEdgeListOffset() uint64 {
+	// Edge list comes right after document content and before doc stored offsets
+	// The stored index offset points to where the doc offsets start
+	// So edge list starts right after the last document content
+	// which is at sb.storedIndexOffset + (8 * sb.numDocs)
+	// since each doc offset is 8 bytes
+	return sb.storedIndexOffset + (8 * sb.numDocs)
+}

section_inverted_text_index.go

@@ -1046,6 +1046,7 @@ func (io *invertedIndexOpaque) Reset() (err error) {
 	atomic.StoreUint64(&io.bytesWritten, 0)
 	io.fieldsSame = false
 	io.numDocs = 0
+	io.fieldAddrs = map[int]int{}
 
 	return err
 }

section_synonym_index.go

@@ -137,6 +137,7 @@ func (so *synonymIndexOpaque) Reset() (err error) {
 	so.FieldIDtoThesaurusID = nil
 	so.SynonymTermToID = so.SynonymTermToID[:0]
 	so.SynonymIDtoTerm = so.SynonymIDtoTerm[:0]
+	so.thesaurusAddrs = map[int]int{}
 
 	so.tmp0 = so.tmp0[:0]
 	return err