Add optimal partitioning

Author: ekzhu

README.md

@@ -92,25 +92,48 @@ package.
 sort.Sort(lshensemble.BySize(domainRecords))
 ```
 
-Now you can use `BootstrapLshEnsemble` 
-(or `BootstrapLshEnsemblePlus` for better accuracy at higher memory cost\*) 
-to create an LSH Ensemble index. You need to
+Now you can use `BootstrapLshEnsembleOptimal`/`BootstrapLshEnsembleEquiDepth`
+(or `BootstrapLshEnsemblePlusOptimal`/`BootstrapLshEnsemblePlusEquiDepth`) 
+for better accuracy at higher memory cost\*) 
+to create an LSH Ensemble index. 
+`BootstrapLshEnsembleOptimal` uses dynamic programming to create partitions that
+are optimal in the sense that the total number of false positives generated from
+all the partitions are minimized. This method can be
+a bit slower due to the dynamic programming overhead, however, it creates
+optimized partitions for any kind of data distribution.
+`BootstrapLshEnsembleEquiDepth` uses simple equi-depth -- same number of domains
+in every partition. This is method is described in the original 
+[paper](http://www.vldb.org/pvldb/vol9/p1185-zhu.pdf) as suitable for power-law
+distributed domain sizes, which is common in real-world domains.
+You need to
 specify the number of partitions to use and some other parameters.
 The LSH parameter K (number of hash functions per band) is dynamically tuned at query-time,
 but the maximum value should be specified here.
 
-\* See [explanation](#maxk-explanation) for the difference.
+\* See [explanation](#maxk-explanation) for the reason for the "Plus" version.
 
 ```go
-// set the number of partitions
+// Set the number of partitions
 numPart := 8
 
-// set the maximum value for the MinHash LSH parameter K 
+// Set the maximum value for the MinHash LSH parameter K 
 // (number of hash functions per band).
 maxK := 4
 
-// create index, you can also use BootstrapLshEnsemblePlus for better accuracy
-index, err := lshensemble.BootstrapLshEnsemble(numPart, numHash, maxK, len(domainRecords), lshensemble.Recs2Chan(domainRecords))
+// Create index using equi-depth partitioning
+// You can also use BootstrapLshEnsemblePlusEquiDepth for better accuracy
+index_eqd, err := lshensemble.BootstrapLshEnsembleEquiDepth(numPart, numHash, maxK, 
+    len(domainRecords), lshensemble.Recs2Chan(domainRecords))
+if err != nil {
+	panic(err)
+}
+
+// Create index using optimal partitioning
+// You can also use BootstrapLshEnsemblePlusOptimal for better accuracy
+index_opt, err := lshensemble.BootstrapLshEnsembleOptimal(numPart, numHash, maxK,
+    func () <-chan *lshensemble.DomainRecord { 
+        return lshensemble.Recs2Chan(domainRecords); 
+    })
 if err != nil {
 	panic(err)
 }
@@ -119,7 +142,7 @@ if err != nil {
 For better memory efficiency when the number of domains is large, 
 it's wiser to use Golang channels and goroutines
 to pipeline the generation of the signatures, and then use disk-based sorting to sort the domain records. 
-This is why `BootstrapLshEnsemble` accepts a channel of `*DomainRecord` as input.
+This is why `BootstrapLshEnsembleEquiDepth` accepts a channel of `*DomainRecord` as input.
 For a small number of domains, you simply use `Recs2Chan` to convert the sorted slice of `*DomainRecord`
 into a `chan *DomainRecord`.
 To help serializing the domain records to disk, you can use `SerializeSignature`
@@ -223,9 +246,11 @@ Essentially, we have less freedom in varying `L`, as
 In this library for LSH Ensemble, we provide both implmentations 
 (LSH Forest and "vanilla" MinHash LSH ).
 Specifically, 
-* `BootstrapLshEnsemble` builds the index using the LSH Forest implementation, 
+* `BootstrapLshEnsembleEquiDepth` and `BootstrapLshEnsembleOptimal` 
+build the index using the LSH Forest implementation, 
 which use less memory but with a more restricted parameter space for optimization.
-* `BootstrapLshEnsemblePlus` builds the index using the "vanilla" MinHash LSH
+* `BootstrapLshEnsemblePlusEquiDepth` and `BootstrapLshEnsemblePlusOptimal` 
+build the index using the "vanilla" MinHash LSH
 implementation (one LSH for every `K`), which uses more memory (bounded by `MaxK`)
 but with no restriction on `L`.
 

accuracy_benchmark_test.go

@@ -10,7 +10,7 @@ import (
 	"time"
 )
 
-func benchmark_accuracy(groundTruthFilename, queryResultFilename, outputFilename string) {
+func benchmarkAccuracy(groundTruthFilename, queryResultFilename, outputFilename string) {
 	groundTruths := readQueryResultFile(groundTruthFilename)
 	queryResults := readQueryResultFile(queryResultFilename)
 	precisions := make([]float64, 0)
@@ -67,7 +67,7 @@ func recallPrecision(result, groundTruth queryResult) (recall, precision float64
 	overlap := 0
 	for id := range test {
 		if _, found := truth[id]; found {
-			overlap += 1
+			overlap++
 		}
 	}
 	recall = float64(overlap) / float64(len(truth))
@@ -82,6 +82,7 @@ func readQueryResultFile(queryResultFile string) []queryResult {
 		panic(err)
 	}
 	scanner := bufio.NewScanner(file)
+	scanner.Buffer(nil, 4096*1024*1024)
 	for scanner.Scan() {
 		raw := strings.Split(scanner.Text(), "\t")
 		key := raw[0]

bootstrap.go

@@ -3,17 +3,85 @@ package lshensemble
 import "errors"
 
 var (
-	ErrDomainSizeOrder = errors.New("Domain records must be sorted in ascending order of size")
+	errDomainSizeOrder = errors.New("Domain records must be sorted in ascending order of size")
 )
 
-func bootstrap(index *LshEnsemble, totalNumDomains int, sortedDomains <-chan *DomainRecord) error {
+func bootstrapOptimalPartitions(domains <-chan *DomainRecord, numPart int) []Partition {
+	sizes, counts := computeSizeDistribution(domains)
+	partitions := optimalPartitions(sizes, counts, numPart)
+	return partitions
+}
+
+func bootstrapOptimal(index *LshEnsemble, sortedDomains <-chan *DomainRecord) error {
+	var currPart int
+	var currSize int
+	for rec := range sortedDomains {
+		if currSize > rec.Size {
+			return errDomainSizeOrder
+		}
+		currSize = rec.Size
+		if currSize > index.Partitions[currPart].Upper {
+			currPart++
+		}
+		if currPart >= len(index.Partitions) ||
+			!(index.Partitions[currPart].Lower <= currSize &&
+				currSize <= index.Partitions[currPart].Upper) {
+			return errors.New("Domain records does not match the existing partitions")
+		}
+		index.Add(rec.Key, rec.Signature, currPart)
+	}
+	index.Index()
+	return nil
+}
+
+// BootstrapLshEnsembleOptimal builds an index from domains using optimal
+// partitioning.
+// The returned index consists of MinHash LSH implemented using LshForest.
+// numPart is the number of partitions to create.
+// numHash is the number of hash functions in MinHash.
+// maxK is the maximum value for the MinHash parameter K - the number of hash
+// functions per "band".
+// sortedDomainFactory is factory function that returns a DomainRecord channel
+// emitting domains in sorted order by their sizes.
+func BootstrapLshEnsembleOptimal(numPart, numHash, maxK int,
+	sortedDomainFactory func() <-chan *DomainRecord) (*LshEnsemble, error) {
+	partitions := bootstrapOptimalPartitions(sortedDomainFactory(), numPart)
+	index := NewLshEnsemble(partitions, numHash, maxK)
+	err := bootstrapOptimal(index, sortedDomainFactory())
+	if err != nil {
+		return nil, err
+	}
+	return index, nil
+}
+
+// BootstrapLshEnsemblePlusOptimal builds an index from domains using optimal
+// partitioning.
+// The returned index consists of MinHash LSH implemented using LshForestArray.
+// numPart is the number of partitions to create.
+// numHash is the number of hash functions in MinHash.
+// maxK is the maximum value for the MinHash parameter K - the number of hash
+// functions per "band".
+// sortedDomainFactory is factory function that returns a DomainRecord channel
+// emitting domains in sorted order by their sizes.
+func BootstrapLshEnsemblePlusOptimal(numPart, numHash, maxK int,
+	sortedDomainFactory func() <-chan *DomainRecord) (*LshEnsemble, error) {
+	partitions := bootstrapOptimalPartitions(sortedDomainFactory(), numPart)
+	index := NewLshEnsemblePlus(partitions, numHash, maxK)
+	err := bootstrapOptimal(index, sortedDomainFactory())
+	if err != nil {
+		return nil, err
+	}
+	return index, nil
+}
+
+func bootstrapEquiDepth(index *LshEnsemble, totalNumDomains int, sortedDomains <-chan *DomainRecord) error {
 	numPart := len(index.Partitions)
 	depth := totalNumDomains / numPart
 	var currDepth, currPart int
 	var currSize int
 	for rec := range sortedDomains {
 		if currSize > rec.Size {
-			return ErrDomainSizeOrder
+			return errDomainSizeOrder
 		}
 		currSize = rec.Size
 		index.Add(rec.Key, rec.Signature, currPart)
@@ -29,30 +97,36 @@ func bootstrap(index *LshEnsemble, totalNumDomains int, sortedDomains <-chan *Do
 	return nil
 }
 
-// BoostrapLshEnsemble builds an index from a channel of domains.
+// BootstrapLshEnsembleEquiDepth builds an index from a channel of domains
+// using equi-depth partitions -- partitions have approximately the same
+// number of domains.
 // The returned index consists of MinHash LSH implemented using LshForest.
 // numPart is the number of partitions to create.
 // numHash is the number of hash functions in MinHash.
 // maxK is the maximum value for the MinHash parameter K - the number of hash functions per "band".
 // sortedDomains is a DomainRecord channel emitting domains in sorted order by their sizes.
-func BootstrapLshEnsemble(numPart, numHash, maxK, totalNumDomains int, sortedDomains <-chan *DomainRecord) (*LshEnsemble, error) {
+func BootstrapLshEnsembleEquiDepth(numPart, numHash, maxK, totalNumDomains int,
+	sortedDomains <-chan *DomainRecord) (*LshEnsemble, error) {
 	index := NewLshEnsemble(make([]Partition, numPart), numHash, maxK)
-	err := bootstrap(index, totalNumDomains, sortedDomains)
+	err := bootstrapEquiDepth(index, totalNumDomains, sortedDomains)
 	if err != nil {
 		return nil, err
 	}
 	return index, nil
 }
 
-// BoostrapLshEnsemblePlus builds an index from a channel of domains.
+// BootstrapLshEnsemblePlusEquiDepth builds an index from a channel of domains
+// using equi-depth partitions -- partitions have approximately the same
+// number of domains.
 // The returned index consists of MinHash LSH implemented using LshForestArray.
 // numPart is the number of partitions to create.
 // numHash is the number of hash functions in MinHash.
 // maxK is the maximum value for the MinHash parameter K - the number of hash functions per "band".
 // sortedDomains is a DomainRecord channel emitting domains in sorted order by their sizes.
-func BootstrapLshEnsemblePlus(numPart, numHash, maxK, totalNumDomains int, sortedDomains <-chan *DomainRecord) (*LshEnsemble, error) {
+func BootstrapLshEnsemblePlusEquiDepth(numPart, numHash, maxK,
+	totalNumDomains int, sortedDomains <-chan *DomainRecord) (*LshEnsemble, error) {
 	index := NewLshEnsemblePlus(make([]Partition, numPart), numHash, maxK)
-	err := bootstrap(index, totalNumDomains, sortedDomains)
+	err := bootstrapEquiDepth(index, totalNumDomains, sortedDomains)
 	if err != nil {
 		return nil, err
 	}

cod_benchmark_test.go

@@ -13,9 +13,15 @@ import (
 	"time"
 )
 
-const benchmarkSeed = 42
-const fracQuery = 0.01
-const minDomainSize = 10
+const (
+	benchmarkSeed = 42
+	fracQuery     = 0.01
+	minDomainSize = 10
+)
+
+var (
+	thresholds = []float64{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0}
+)
 
 // Running this function requires a `_cod_domains` directory
 // in the current directory.
@@ -49,17 +55,19 @@ func Benchmark_CanadianOpenData(b *testing.B) {
 	}
 
 	// Run benchmark
-	log.Printf("Canadian Open Data benchmark threshold = %.2f", 0.5)
-	benchmark_cod(rawDomains, queries, 0.5)
+	for _, threshold := range thresholds {
+		log.Printf("Canadian Open Data benchmark threshold = %.2f", threshold)
+		benchmarkCOD(rawDomains, queries, threshold)
+	}
 }
 
-func benchmark_cod(rawDomains, queries []rawDomain, threshold float64) {
-	linearscan_output := fmt.Sprintf("_cod_linearscan_threshold=%.2f", threshold)
-	lshensemble_output := fmt.Sprintf("_cod_lshensemble_threshold=%.2f", threshold)
-	accuracy_output := fmt.Sprintf("_cod_accuracy_threhsold=%.2f", threshold)
-	benchmark_linearscan(rawDomains, queries, threshold, linearscan_output)
-	benchmark_lshensemble(rawDomains, queries, threshold, lshensemble_output)
-	benchmark_accuracy(linearscan_output, lshensemble_output, accuracy_output)
+func benchmarkCOD(rawDomains, queries []rawDomain, threshold float64) {
+	linearscanOutput := fmt.Sprintf("_cod_linearscan_threshold=%.2f", threshold)
+	lshensembleOutput := fmt.Sprintf("_cod_lshensemble_threshold=%.2f", threshold)
+	accuracyOutput := fmt.Sprintf("_cod_accuracy_threshold=%.2f", threshold)
+	benchmarkLinearscan(rawDomains, queries, threshold, linearscanOutput)
+	benchmarkLshEnsemble(rawDomains, queries, threshold, lshensembleOutput)
+	benchmarkAccuracy(linearscanOutput, lshensembleOutput, accuracyOutput)
 }
 
 type rawDomain struct {

linearscan_benchmark_test.go

@@ -5,7 +5,8 @@ import (
 	"time"
 )
 
-func benchmark_linearscan(rawDomains []rawDomain, queries []rawDomain, threshold float64, outputFilename string) {
+func benchmarkLinearscan(rawDomains []rawDomain, queries []rawDomain,
+	threshold float64, outputFilename string) {
 	log.Printf("Start Linear Scan with %d queries", len(queries))
 	results := make(chan queryResult)
 	go func() {

lshensemble.go

@@ -1,6 +1,7 @@
 package lshensemble
 
 import (
+	"errors"
 	"fmt"
 	"sync"
 	"time"
@@ -87,6 +88,19 @@ func (e *LshEnsemble) Add(key interface{}, sig []uint64, partInd int) {
 	e.lshes[partInd].Add(key, sig)
 }
 
+// Prepare adds a new domain to the index given its size, and partition will
+// be selected automatically. It could be more efficient to use Add().
+// The added domain won't be searchable until the Index() function is called.
+func (e *LshEnsemble) Prepare(key interface{}, sig []uint64, size int) error {
+	for i := range e.Partitions {
+		if size >= e.Partitions[i].Lower && size <= e.Partitions[i].Upper {
+			e.Add(key, sig, i)
+			break
+		}
+	}
+	return errors.New("No matching partition found")
+}
+
 // Index makes all added domains searchable.
 func (e *LshEnsemble) Index() {
 	for i := range e.lshes {

lshensemble_benchmark_test.go

@@ -6,7 +6,16 @@ import (
 	"time"
 )
 
-func benchmark_lshensemble(rawDomains []rawDomain, rawQueries []rawDomain, threshold float64, outputFilename string) {
+const (
+	numHash = 256
+	numPart = 32
+	maxK    = 4
+	// useOptimalPartitions = true
+	useOptimalPartitions = false
+)
+
+func benchmarkLshEnsemble(rawDomains []rawDomain, rawQueries []rawDomain,
+	threshold float64, outputFilename string) {
 	numHash := 256
 	numPart := 32
 	maxK := 4
@@ -27,8 +36,14 @@ func benchmark_lshensemble(rawDomains []rawDomain, rawQueries []rawDomain, thres
 	// Indexing
 	log.Print("Start building LSH Ensemble index")
 	sort.Sort(BySize(domainRecords))
-	index, _ := BootstrapLshEnsemblePlus(numPart, numHash, maxK, len(domainRecords),
-		Recs2Chan(domainRecords))
+	var index *LshEnsemble
+	if useOptimalPartitions {
+		index, _ = BootstrapLshEnsemblePlusOptimal(numPart, numHash, maxK,
+			func() <-chan *DomainRecord { return Recs2Chan(domainRecords) })
+	} else {
+		index, _ = BootstrapLshEnsemblePlusEquiDepth(numPart, numHash, maxK,
+			len(domainRecords), Recs2Chan(domainRecords))
+	}
 	log.Print("Finished building LSH Ensemble index")
 	// Querying
 	log.Printf("Start querying LSH Ensemble index with %d queries", len(queries))