Refactor sequence processing for parallel loading and memory management

This commit introduces significant changes to the sequence processing pipeline, implementing parallel loading, k-mer budget control, explicit memory management, and refined logic for pair evaluation and resolution.

pkg/obikmer/seqquorum.go

@@ -24,9 +24,12 @@ import (
 //  2. CountAtLeast processes the partitions in parallel: each one is loaded,
 //     expanded to (canonical k-mer, sequence id) pairs, sorted, and reduced.
 //
-// Memory during phase 2 is about 16 bytes per k-mer occurrence of the largest
-// partitions being processed at the same time, so P must grow with the total
-// input size (total bases / P is the occurrence count of an average partition).
+// Memory during phase 2 is 16 bytes per k-mer occurrence of the partitions
+// being processed at the same time. Their number is bounded by an occurrence
+// budget (see SetOccurrenceBudget), so that the memory does not depend on how
+// the k-mers happen to be spread over the partitions; P only has to be large
+// enough for the largest partition to be small against the budget (total bases
+// / P is the occurrence count of an average partition, see DefaultPartitions).
 //
 // Bases other than A, C, G, T, U break a sequence into independent runs: no
 // k-mer ever contains an ambiguous base.
@@ -39,13 +42,16 @@ type SeqKmerCounter struct {
 	mu       sync.Mutex
 	firstErr error
 	sealed   bool
+
+	budget int // k-mer occurrences held at once in phase 2
 }
 
 type seqPartWriter struct {
-	mu   sync.Mutex
-	w    *bufio.Writer
-	file *os.File
-	path string
+	mu    sync.Mutex
+	w     *bufio.Writer
+	file  *os.File
+	path  string
+	kmers int // k-mer occurrences written, so that the partition is loaded in an array of the exact size
 }
 
 // SeqKmerCount is the result for one k-mer.
@@ -78,7 +84,7 @@ func NewSeqKmerCounter(parentDir string, k, m, partitions int) (*SeqKmerCounter,
 		return nil, err
 	}
 
-	c := &SeqKmerCounter{dir: dir, k: k, m: m, p: partitions}
+	c := &SeqKmerCounter{dir: dir, k: k, m: m, p: partitions, budget: DefaultOccurrenceBudget}
 	c.parts = make([]*seqPartWriter, partitions)
 	for i := range c.parts {
 		path := filepath.Join(dir, fmt.Sprintf("part_%04d.sqk", i))
@@ -92,6 +98,38 @@ func NewSeqKmerCounter(parentDir string, k, m, partitions int) (*SeqKmerCounter,
 	return c, nil
 }
 
+// DefaultOccurrenceBudget is the number of k-mer occurrences phase 2 holds at
+// once by default: 32 million, 512 MB. A partition larger than the budget is
+// processed alone. Measured on 16 cores, 526 million occurrences to count: with
+// 32M the count is as fast as without any limit; 16M make it 15% slower, 8M
+// 60% slower.
+const DefaultOccurrenceBudget = 32 << 20
+
+// SetOccurrenceBudget bounds the k-mer occurrences (16 bytes each) that
+// CountAtLeast holds in memory at once. It must be called before CountAtLeast.
+func (c *SeqKmerCounter) SetOccurrenceBudget(occurrences int) {
+	if occurrences > 0 {
+		c.budget = occurrences
+	}
+}
+
+// partitionOccurrences is the occurrence count aimed at for an average
+// partition by DefaultPartitions (a partition of 1M occurrences is 16 MB).
+const partitionOccurrences = 1 << 20
+
+// DefaultPartitions returns a partition count suited to a total of kmers
+// k-mer occurrences: 256 at least, and more, as a power of two up to 1024, when
+// the average partition would otherwise exceed about a million occurrences.
+// The largest partitions hold several times the average. The limit keeps the
+// number of simultaneously open files reasonable.
+func DefaultPartitions(kmers int) int {
+	p := 256
+	for p < 1024 && kmers/p > partitionOccurrences {
+		p *= 2
+	}
+	return p
+}
+
 // K returns the k-mer size.
 func (c *SeqKmerCounter) K() int { return c.k }
 
@@ -148,6 +186,7 @@ func (c *SeqKmerCounter) addRun(id uint32, run []byte, scratch []byte) []byte {
 		pw := c.parts[int(sk.Minimizer%uint64(c.p))]
 		pw.mu.Lock()
 		_, err := pw.w.Write(scratch)
+		pw.kmers += len(sk.Sequence) - c.k + 1
 		pw.mu.Unlock()
 		if err != nil {
 			c.setErr(err)