<commit_before>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package recon contains code for recon.
package recon

import (
	"context"
	"fmt"
	"sort"

	pb "github.com/datacommonsorg/mixer/internal/proto"
	"github.com/datacommonsorg/mixer/internal/server/v1/propertyvalues"
	"github.com/datacommonsorg/mixer/internal/store"
	"github.com/datacommonsorg/mixer/internal/util"
	"golang.org/x/sync/errgroup"
	"googlemaps.github.io/maps"
)

const (
	maxNumEntitiesPerRequest  = 5000
	maxMapsAPICallsInParallel = 25
)

type entityInfo struct {
	description string
	typ         string
}

// BulkFindEntities implements API for Mixer.BulkFindEntities.
func BulkFindEntities(
	ctx context.Context,
	in *pb.BulkFindEntitiesRequest,
	store *store.Store,
	mapsClient *maps.Client,
) (*pb.BulkFindEntitiesResponse, error) {
	if l := len(in.GetEntities()); l == 0 {
		return nil, fmt.Errorf("empty input")
	} else if l > maxNumEntitiesPerRequest {
		return nil, fmt.Errorf(
			"exceeded max number of entities per request (%d): %d",
			maxNumEntitiesPerRequest, l)
	}

	// Load input.
	entityInfoSet := map[entityInfo]struct{}{}
	for _, entity := range in.GetEntities() {
		description := entity.GetDescription()
		if description == "" {
			continue
		}
		entityInfoSet[entityInfo{description, entity.GetType()}] = struct{}{}
	}

	// Get place IDs.
	entityInfoToPlaceIDs, placeIDSet, err := resolvePlaceIDs(ctx, mapsClient, entityInfoSet)
	if err != nil {
		return nil, err
	}

	// Resolve place IDs to get DCIDs.
	placeIDToDCIDs, dcidSet, err := resolveDCIDs(ctx, placeIDSet, store)
	if err != nil {
		return nil, err
	}

	// Get types of the DCIDs.
	dcidToTypeSet, err := getPlaceTypes(ctx, dcidSet, store)
	if err != nil {
		return nil, err
	}

	// Assemble results.
	resp := &pb.BulkFindEntitiesResponse{}
	for entityInfo, placeIDs := range entityInfoToPlaceIDs {
		entity := &pb.BulkFindEntitiesResponse_Entity{
			Description: entityInfo.description,
			Type:        entityInfo.typ,
		}

		allDCIDs := []string{}
		for _, placeID := range placeIDs {
			if dcids, ok := placeIDToDCIDs[placeID]; ok {
				allDCIDs = append(allDCIDs, dcids...)
			}
		}

		if len(allDCIDs) != 0 {
			if entityInfo.typ == "" {
				// No type filtering.
				entity.Dcids = allDCIDs
			} else {
				// Type filtering.
				filteredDCIDs := []string{}
				for _, dcid := range allDCIDs {
					typeSet, ok := dcidToTypeSet[dcid]
					if !ok {
						continue
					}
					if _, ok := typeSet[entityInfo.typ]; ok {
						filteredDCIDs = append(filteredDCIDs, dcid)
					}
				}
				entity.Dcids = filteredDCIDs
			}
		}

		resp.Entities = append(resp.Entities, entity)
	}

	// Sort to make results determistic.
	sort.Slice(resp.Entities, func(i, j int) bool {
		if resp.Entities[i].GetDescription() == resp.Entities[j].GetDescription() {
			return resp.Entities[i].GetType() < resp.Entities[j].GetType()
		}
		return resp.Entities[i].GetDescription() < resp.Entities[j].GetDescription()
	})

	return resp, nil
}

func resolvePlaceIDs(
	ctx context.Context,
	mapsClient *maps.Client,
	entityInfoSet map[entityInfo]struct{},
) (
	map[entityInfo][]string, /* entityInfo -> [place ID] */
	map[string]struct{}, /* [place ID] for all entities */
	error,
) {
	type resolveResult struct {
		entityInfo *entityInfo
		placeIDs   []string
	}

	// Distribute entityInfoSet to maxMapsAPICallsInParallel shards for parallel processing.
	entityInfoListShards := make([][]entityInfo, maxMapsAPICallsInParallel)
	idx := 0
	for entityInfo := range entityInfoSet {
		entityInfoListShards[idx] = append(entityInfoListShards[idx], entityInfo)
		idx++
		if idx >= maxMapsAPICallsInParallel {
			idx = 0
		}
	}

	// The channel to receive results from parallel workers.
	resolveResultChan := make(chan resolveResult, len(entityInfoSet))

	// Worker function.
	mapsAPICallWorkerFunc := func(ctx context.Context, i int) func() error {
		return func() error {
			for _, entityInfo := range entityInfoListShards[i] {
				placeIDs, err := findPlaceIDsForEntity(ctx, mapsClient, &entityInfo)
				if err != nil {
					return err
				}
				// Only keep the first place ID, as the rest ones are usually much less accurate.
				resolveResultChan <- resolveResult{
					entityInfo: &entityInfo,
					placeIDs:   []string{placeIDs[0]}}
			}
			return nil
		}
	}

	// Call Maps API to find place IDs in parallel.
	// The errors in the Goroutines need to be captured, so we use errgroup.
	eg, errCtx := errgroup.WithContext(ctx)
	for i := 0; i < maxMapsAPICallsInParallel; i++ {
		eg.Go(mapsAPICallWorkerFunc(errCtx, i))
	}
	if err := eg.Wait(); err != nil {
		return nil, nil, err
	}
	close(resolveResultChan)

	// Read out the results sent by workers.
	entityInfoToPlaceIDs := map[entityInfo][]string{}
	placeIDSet := map[string]struct{}{}
	for res := range resolveResultChan {
		entityInfoToPlaceIDs[*res.entityInfo] = res.placeIDs
		for _, placeID := range res.placeIDs {<commit_msg>Fix failing recon API test (#976)<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package recon contains code for recon.
package recon

import (
	"context"
	"fmt"
	"sort"

	pb "github.com/datacommonsorg/mixer/internal/proto"
	"github.com/datacommonsorg/mixer/internal/server/v1/propertyvalues"
	"github.com/datacommonsorg/mixer/internal/store"
	"github.com/datacommonsorg/mixer/internal/util"
	"golang.org/x/sync/errgroup"
	"googlemaps.github.io/maps"
)

const (
	maxNumEntitiesPerRequest  = 5000
	maxMapsAPICallsInParallel = 25
)

type entityInfo struct {
	description string
	typ         string
}

// BulkFindEntities implements API for Mixer.BulkFindEntities.
func BulkFindEntities(
	ctx context.Context,
	in *pb.BulkFindEntitiesRequest,
	store *store.Store,
	mapsClient *maps.Client,
) (*pb.BulkFindEntitiesResponse, error) {
	if l := len(in.GetEntities()); l == 0 {
		return nil, fmt.Errorf("empty input")
	} else if l > maxNumEntitiesPerRequest {
		return nil, fmt.Errorf(
			"exceeded max number of entities per request (%d): %d",
			maxNumEntitiesPerRequest, l)
	}

	// Load input.
	entityInfoSet := map[entityInfo]struct{}{}
	for _, entity := range in.GetEntities() {
		description := entity.GetDescription()
		if description == "" {
			continue
		}
		entityInfoSet[entityInfo{description, entity.GetType()}] = struct{}{}
	}

	// Get place IDs.
	entityInfoToPlaceIDs, placeIDSet, err := resolvePlaceIDs(ctx, mapsClient, entityInfoSet)
	if err != nil {
		return nil, err
	}

	// Resolve place IDs to get DCIDs.
	placeIDToDCIDs, dcidSet, err := resolveDCIDs(ctx, placeIDSet, store)
	if err != nil {
		return nil, err
	}

	// Get types of the DCIDs.
	dcidToTypeSet, err := getPlaceTypes(ctx, dcidSet, store)
	if err != nil {
		return nil, err
	}

	// Assemble results.
	resp := &pb.BulkFindEntitiesResponse{}
	for entityInfo, placeIDs := range entityInfoToPlaceIDs {
		entity := &pb.BulkFindEntitiesResponse_Entity{
			Description: entityInfo.description,
			Type:        entityInfo.typ,
		}

		allDCIDs := []string{}
		for _, placeID := range placeIDs {
			if dcids, ok := placeIDToDCIDs[placeID]; ok {
				allDCIDs = append(allDCIDs, dcids...)
			}
		}

		if len(allDCIDs) != 0 {
			if entityInfo.typ == "" {
				// No type filtering.
				entity.Dcids = allDCIDs
			} else {
				// Type filtering.
				filteredDCIDs := []string{}
				for _, dcid := range allDCIDs {
					typeSet, ok := dcidToTypeSet[dcid]
					if !ok {
						continue
					}
					if _, ok := typeSet[entityInfo.typ]; ok {
						filteredDCIDs = append(filteredDCIDs, dcid)
					}
				}
				entity.Dcids = filteredDCIDs
			}
		}

		resp.Entities = append(resp.Entities, entity)
	}

	// Sort to make results determistic.
	sort.Slice(resp.Entities, func(i, j int) bool {
		if resp.Entities[i].GetDescription() == resp.Entities[j].GetDescription() {
			return resp.Entities[i].GetType() < resp.Entities[j].GetType()
		}
		return resp.Entities[i].GetDescription() < resp.Entities[j].GetDescription()
	})

	return resp, nil
}

func resolvePlaceIDs(
	ctx context.Context,
	mapsClient *maps.Client,
	entityInfoSet map[entityInfo]struct{},
) (
	map[entityInfo][]string, /* entityInfo -> [place ID] */
	map[string]struct{}, /* [place ID] for all entities */
	error,
) {
	type resolveResult struct {
		entityInfo *entityInfo
		placeIDs   []string
	}

	// Distribute entityInfoSet to maxMapsAPICallsInParallel shards for parallel processing.
	entityInfoListShards := make([][]entityInfo, maxMapsAPICallsInParallel)
	idx := 0
	for entityInfo := range entityInfoSet {
		entityInfoListShards[idx] = append(entityInfoListShards[idx], entityInfo)
		idx++
		if idx >= maxMapsAPICallsInParallel {
			idx = 0
		}
	}

	// The channel to receive results from parallel workers.
	resolveResultChan := make(chan resolveResult, len(entityInfoSet))

	// Worker function.
	mapsAPICallWorkerFunc := func(ctx context.Context, i int) func() error {
		return func() error {
			for _, entityInfo := range entityInfoListShards[i] {
				placeIDs, err := findPlaceIDsForEntity(ctx, mapsClient, &entityInfo)
				if err != nil {
					return err
				}
				usedPlaceIds := []string{}
				if len(placeIDs) > 0 {
					// Only keep the first place ID, as the rest ones are usually much less accurate.
					usedPlaceIds = []string{placeIDs[0]}
				}
				resolveResultChan <- resolveResult{
					entityInfo: &entityInfo,
					placeIDs:   usedPlaceIds}
			}
			return nil
		}
	}

	// Call Maps API to find place IDs in parallel.
	// The errors in the Goroutines need to be captured, so we use errgroup.
	eg, errCtx := errgroup.WithContext(ctx)
	for i := 0; i < maxMapsAPICallsInParallel; i++ {
		eg.Go(mapsAPICallWorkerFunc(errCtx, i))
	}
	if err := eg.Wait(); err != nil {
		return nil, nil, err
	}
	close(resolveResultChan)

	// Read out the results sent by workers.
	entityInfoToPlaceIDs := map[entityInfo][]string{}
	placeIDSet := map[string]struct{}{}
	for res := range resolveResultChan {
		entityInfoToPlaceIDs[*res.entityInfo] = res.placeIDs
		for _, placeID := range res.placeIDs {
{'fraction_non_alphanumeric': 0.11865482233502538, 'fraction_numerical': 0.004169688179840464, 'mean_word_length': 3.675, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 10, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2634133', 'n_tokens_mistral': 3922, 'n_tokens_neox': 3591, 'n_words': 988}