Add tag and face counts to the photo index (#2660)

samples/go/photo-index/main.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"sync"
 
 	"github.com/attic-labs/noms/go/datas"
 	"github.com/attic-labs/noms/go/spec"
@@ -93,6 +94,10 @@ func index() (win bool) {
 	byTag := types.NewGraphBuilder(db, types.MapKind, true)
 	byFace := types.NewGraphBuilder(db, types.MapKind, true)
 
+	tagCounts := map[types.String]int{}
+	faceCounts := map[types.String]int{}
+	countsMtx := sync.Mutex{}
+
 	for _, v := range inputs {
 		walk.WalkValues(v, db, func(cv types.Value) (stop bool) {
 			if types.IsSubtype(photoType, cv.Type()) {
@@ -113,17 +118,31 @@ func index() (win bool) {
 				byDate.SetInsert([]types.Value{d}, cv)
 
 				// Index by tag, then date
+				moreTags := map[types.String]int{}
 				s.Get("tags").(types.Set).IterAll(func(t types.Value) {
 					byTag.SetInsert([]types.Value{t, d}, cv)
+					moreTags[t.(types.String)]++
 				})
 
 				// Index by face, then date
+				moreFaces := map[types.String]int{}
 				if types.IsSubtype(withFaces, cv.Type()) {
 					s.Get("faces").(types.Set).IterAll(func(t types.Value) {
-						byFace.SetInsert([]types.Value{t.(types.Struct).Get("name"), d}, cv)
+						name := t.(types.Struct).Get("name").(types.String)
+						byFace.SetInsert([]types.Value{name, d}, cv)
+						moreFaces[name]++
 					})
 				}
 
+				countsMtx.Lock()
+				for tag, count := range moreTags {
+					tagCounts[tag] += count
+				}
+				for face, count := range moreFaces {
+					faceCounts[face] += count
+				}
+				countsMtx.Unlock()
+
 				// Can't be any photos inside photos, so we can save a little bit here.
 				stop = true
 			}
@@ -132,9 +151,11 @@ func index() (win bool) {
 	}
 
 	outDS, err = db.CommitValue(outDS, types.NewStruct("", types.StructData{
-		"byDate": byDate.Build(),
-		"byTag":  byTag.Build(),
-		"byFace": byFace.Build(),
+		"byDate":       byDate.Build(),
+		"byTag":        byTag.Build(),
+		"byFace":       byFace.Build(),
+		"tagsByCount":  stringsByCount(db, tagCounts),
+		"facesByCount": stringsByCount(db, faceCounts),
 	}))
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Could not commit: %s\n", err)
@@ -145,6 +166,15 @@ func index() (win bool) {
 	return
 }
 
+func stringsByCount(db datas.Database, strings map[types.String]int) types.Map {
+	b := types.NewGraphBuilder(db, types.MapKind, true)
+	for s, count := range strings {
+		// Sort by largest count by negating.
+		b.SetInsert([]types.Value{types.Number(-count)}, s)
+	}
+	return b.Build().(types.Map)
+}
+
 func usage() {
 	fmt.Fprintf(os.Stderr, "photo-index indexes photos by common attributes\n\n")
 	fmt.Fprintf(os.Stderr, "Usage: %s -db=<db-spec> -out-ds=<name> [input-paths...]\n\n", path.Base(os.Args[0]))

samples/go/photo-index/main_test.go

@@ -27,6 +27,11 @@ func (s *testSuite) TestWin() {
 	sp := fmt.Sprintf("ldb:%s::test", s.LdbDir)
 	db, ds, _ := spec.GetDataset(sp)
 
+	type Face struct {
+		Name       string
+		X, Y, W, H int
+	}
+
 	type Date struct {
 		NsSinceEpoch int
 	}
@@ -34,6 +39,7 @@ func (s *testSuite) TestWin() {
 	type Photo struct {
 		Title string
 		Tags  types.Set
+		Faces types.Set
 		Sizes map[struct {
 			Width  int
 			Height int
@@ -51,6 +57,19 @@ func (s *testSuite) TestWin() {
 		return s
 	}
 
+	getFaces := func(n int) types.Set {
+		set := types.NewSet()
+		for i := 0; i < n; i++ {
+			v, err := marshal.Marshal(Face{
+				fmt.Sprintf("harry%d", i),
+				i, i, n, n,
+			})
+			s.NoError(err)
+			set = set.Insert(v)
+		}
+		return set
+	}
+
 	getPhoto := func(n int) Photo {
 		return Photo{
 			Title: fmt.Sprintf("photo %d", n),
@@ -60,6 +79,7 @@ func (s *testSuite) TestWin() {
 			DateTaken:     Date{n * 10},
 			DatePublished: Date{n*10 + 1},
 			DateUpdated:   Date{n*10 + 2},
+			Faces:         getFaces(n),
 		}
 	}
 
@@ -78,8 +98,11 @@ func (s *testSuite) TestWin() {
 
 	db, ds, _ = spec.GetDataset(fmt.Sprintf("%s::idx", s.LdbDir))
 	var idx struct {
-		ByDate map[int]types.Set
-		ByTag  map[string]map[int]types.Set
+		ByDate       map[int]types.Set
+		ByTag        map[string]map[int]types.Set
+		ByFace       map[string]map[int]types.Set
+		TagsByCount  map[int]types.Set
+		FacesByCount map[int]types.Set
 	}
 	marshal.Unmarshal(ds.HeadValue(), &idx)
 
@@ -91,9 +114,32 @@ func (s *testSuite) TestWin() {
 	}
 
 	s.Equal(4, len(idx.ByTag))
-	for i := 1; i < 5; i++ {
+	for i := 0; i < 4; i++ {
 		k := fmt.Sprintf("tag%d", i)
 		v := idx.ByTag[k]
 		s.Equal(4-i, len(v))
 	}
+
+	s.Equal(4, len(idx.ByFace))
+	for i := 0; i < 4; i++ {
+		k := fmt.Sprintf("harry%d", i)
+		v := idx.ByFace[k]
+		s.Equal(4-i, len(v))
+	}
+
+	s.Equal(4, len(idx.TagsByCount))
+	for i := 0; i < 4; i++ {
+		tags := idx.TagsByCount[-4+i]
+		s.Equal(1, int(tags.Len()))
+		k := fmt.Sprintf("tag%d", i)
+		s.True(tags.Has(types.String(k)))
+	}
+
+	s.Equal(4, len(idx.FacesByCount))
+	for i := 0; i < 4; i++ {
+		tags := idx.FacesByCount[-4+i]
+		s.Equal(1, int(tags.Len()))
+		k := fmt.Sprintf("harry%d", i)
+		s.True(tags.Has(types.String(k)))
+	}
 }