diff --git a/search.go b/search.go index e324262e..86ea4193 100644 --- a/search.go +++ b/search.go @@ -22,8 +22,10 @@ import ( "github.com/blevesearch/bleve/analysis" "github.com/blevesearch/bleve/analysis/datetime/optional" + "github.com/blevesearch/bleve/document" "github.com/blevesearch/bleve/registry" "github.com/blevesearch/bleve/search" + "github.com/blevesearch/bleve/search/collector" "github.com/blevesearch/bleve/search/query" "github.com/blevesearch/bleve/size" ) @@ -518,3 +520,44 @@ func (sr *SearchResult) Merge(other *SearchResult) { sr.Facets.Merge(other.Facets) } + +// MemoryNeededForSearchResult is an exported helper function to determine the RAM +// needed to accommodate the results for a given search request. +func MemoryNeededForSearchResult(req *SearchRequest) uint64 { + if req == nil { + return 0 + } + + numDocMatches := req.Size + req.From + if req.Size+req.From > collector.PreAllocSizeSkipCap { + numDocMatches = collector.PreAllocSizeSkipCap + } + + estimate := 0 + + // overhead from the SearchResult structure + var sr SearchResult + estimate += sr.Size() + + var dm search.DocumentMatch + sizeOfDocumentMatch := dm.Size() + + // overhead from results + estimate += numDocMatches * sizeOfDocumentMatch + + // overhead from facet results + if req.Facets != nil { + var fr search.FacetResult + estimate += len(req.Facets) * fr.Size() + } + + // highlighting, store + var d document.Document + if len(req.Fields) > 0 || req.Highlight != nil { + for i := 0; i < (req.Size + req.From); i++ { + estimate += (req.Size + req.From) * d.Size() + } + } + + return uint64(estimate) +} diff --git a/search_test.go b/search_test.go index 24249413..87a71828 100644 --- a/search_test.go +++ b/search_test.go @@ -399,3 +399,18 @@ func TestSearchResultFacetsMerge(t *testing.T) { t.Errorf("expected %#v, got %#v", expected, l) } } + +func TestMemoryNeededForSearchResult(t *testing.T) { + query := NewTermQuery("blah") + req := NewSearchRequest(query) + + var sr SearchResult + expect := sr.Size() + var dm search.DocumentMatch + expect += 10 * dm.Size() + + estimate := MemoryNeededForSearchResult(req) + if estimate != uint64(expect) { + t.Errorf("estimate not what is expected: %v != %v", estimate, expect) + } +}