diff --git a/samples/go/photo-index/main.go b/samples/go/photo-index/main.go index fc6cfe33e7..4b975cc03d 100644 --- a/samples/go/photo-index/main.go +++ b/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= -out-ds= [input-paths...]\n\n", path.Base(os.Args[0])) diff --git a/samples/go/photo-index/main_test.go b/samples/go/photo-index/main_test.go index bcd948f350..311b24016d 100644 --- a/samples/go/photo-index/main_test.go +++ b/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))) + } }