added reuse of score arrays and different api for query cells

Author: Likith101

geo/geo_s2plugin_impl.go

@@ -251,6 +251,11 @@ func (p *Point) Cells() ([]uint64, []uint64) {
 	return nil, nil
 }
 
+func (p *Point) QueryCells() ([]uint64, []uint64) {
+	// placeholder implementation
+	return nil, nil
+}
+
 func (p *Point) BoundingBox() index.GeoJSON {
 	// placeholder implementation
 	return nil
@@ -308,6 +313,11 @@ func (br *boundedRectangle) Cells() ([]uint64, []uint64) {
 	return nil, nil
 }
 
+func (br *boundedRectangle) QueryCells() ([]uint64, []uint64) {
+	// placeholder implementation
+	return nil, nil
+}
+
 func (br *boundedRectangle) BoundingBox() index.GeoJSON {
 	// placeholder implementation
 	return nil
@@ -366,6 +376,11 @@ func (bp *boundedPolygon) Cells() ([]uint64, []uint64) {
 	return nil, nil
 }
 
+func (bp *boundedPolygon) QueryCells() ([]uint64, []uint64) {
+	// placeholder implementation
+	return nil, nil
+}
+
 func (bp *boundedPolygon) BoundingBox() index.GeoJSON {
 	// placeholder implementation
 	return nil
@@ -424,6 +439,11 @@ func (pd *pointDistance) Cells() ([]uint64, []uint64) {
 	return nil, nil
 }
 
+func (pd *pointDistance) QueryCells() ([]uint64, []uint64) {
+	// placeholder implementation
+	return nil, nil
+}
+
 func (pd *pointDistance) BoundingBox() index.GeoJSON {
 	// placeholder implementation
 	return nil

geov2/query_contains.go

@@ -34,7 +34,7 @@ type containsQuery struct {
 }
 
 func NewContainsQuery(shape index.GeoJSON) Query {
-	inner, cross := shape.Cells()
+	inner, cross := shape.QueryCells()
 
 	score := CalcCellsScore(inner) + CalcCellsScore(cross)
 
@@ -56,8 +56,12 @@ func (cq *containsQuery) Evaluate(geoData segment.GeoShapeV2Data) *util.Bitset {
 	hits := util.NewBitset(numDocs, exclude)
 	maybeHits := util.NewBitset(numDocs, exclude)
 
-	innerScores := make([]uint64, numDocs)
-	crossScores := make([]uint64, numDocs)
+	// obtain zeroed score arrays from the segment-level pool and return
+	// them once the evaluation is done
+	innerScores := geoData.GetScoreArray()
+	crossScores := geoData.GetScoreArray()
+	defer geoData.PutScoreArray(innerScores)
+	defer geoData.PutScoreArray(crossScores)
 
 	// create an evaluator instance to scan the query cells against the index cells
 	evaluator := NewQueryEvaluator(cq, geoData)

geov2/query_disjoint.go

@@ -29,7 +29,7 @@ type disjointQuery struct {
 }
 
 func NewDisjointQuery(shape index.GeoJSON) Query {
-	inner, cross := shape.Cells()
+	inner, cross := shape.QueryCells()
 
 	return &disjointQuery{
 		innerCells: inner,

geov2/query_intersects.go

@@ -32,7 +32,7 @@ type intersectsQuery struct {
 }
 
 func NewIntersectsQuery(shape index.GeoJSON) Query {
-	inner, cross := shape.Cells()
+	inner, cross := shape.QueryCells()
 
 	return &intersectsQuery{
 		innerCells: inner,
@@ -51,8 +51,12 @@ func (iq *intersectsQuery) Evaluate(geoData segment.GeoShapeV2Data) *util.Bitset
 	hits := util.NewBitset(numDocs, exclude)
 	maybeHits := util.NewBitset(numDocs, exclude)
 
-	innerScores := make([]uint64, numDocs)
-	crossScores := make([]uint64, numDocs)
+	// obtain zeroed score arrays from the segment-level pool and return
+	// them once the evaluation is done
+	innerScores := geoData.GetScoreArray()
+	crossScores := geoData.GetScoreArray()
+	defer geoData.PutScoreArray(innerScores)
+	defer geoData.PutScoreArray(crossScores)
 
 	// create an evaluator instance to scan the query cells against the index cells
 	evaluator := NewQueryEvaluator(iq, geoData)

geov2/query_within.go

@@ -32,7 +32,7 @@ type withinQuery struct {
 }
 
 func NewWithinQuery(shape index.GeoJSON) Query {
-	inner, cross := shape.Cells()
+	inner, cross := shape.QueryCells()
 
 	return &withinQuery{
 		innerCells: inner,
@@ -51,8 +51,13 @@ func (wq *withinQuery) Evaluate(geoData segment.GeoShapeV2Data) *util.Bitset {
 	hits := util.NewBitset(numDocs, exclude)
 	maybeHits := util.NewBitset(numDocs, exclude)
 
-	innerScores := make([]uint64, numDocs)
-	crossScores := make([]uint64, numDocs)
+	// obtain zeroed score arrays from the segment-level pool and return
+	// them once the evaluation is done
+	innerScores := geoData.GetScoreArray()
+	crossScores := geoData.GetScoreArray()
+	defer geoData.PutScoreArray(innerScores)
+	defer geoData.PutScoreArray(crossScores)
+
 	docScores := geoData.DocScores()
 
 	// create an evaluator instance to scan the query cells against the index cells