Optimize TopN() w/ source query.

This commit makes several changes to optimize the TopN() query:

- Reduce highbits() back from 24-bits to 16-bits.
- Reduce MaxArraySize back from 2^20 to 4096.
- Optimize bitmap count invalidation.
- Parallelize TopN() across nodes.
- Parallelize TopN() across slices.

bitmap.go

@@ -316,11 +316,7 @@ func (s *BitmapSegment) ClearBit(i uint64) (changed bool) {
 
 // InvalidateCount updates the cached count in the bitmap.
 func (s *BitmapSegment) InvalidateCount() {
-	itr, n := s.data.Iterator(), uint64(0)
-	for _, eof := itr.Next(); !eof; _, eof = itr.Next() {
-		n++
-	}
-	s.n = n
+	s.n = s.data.Count()
 }
 
 // Bits returns a list of all bits set in the segment.

cmd/pilosactl/main.go

@@ -209,6 +209,12 @@ type ImportCommand struct {
 	// Filenames to import from.
 	Paths []string
 
+	// Size of buffer used to chunk import.
+	BufferSize int
+
+	// Reusable client.
+	Client *pilosa.Client
+
 	// Standard input/output
 	Stdin  io.Reader
 	Stdout io.Writer
@@ -221,6 +227,8 @@ func NewImportCommand(stdin io.Reader, stdout, stderr io.Writer) *ImportCommand
 		Stdin:  stdin,
 		Stdout: stdout,
 		Stderr: stderr,
+
+		BufferSize: 1000000,
 	}
 }
 
@@ -276,41 +284,28 @@ func (cmd *ImportCommand) Run() error {
 	if err != nil {
 		return err
 	}
+	cmd.Client = client
 
 	// Import each path and import by slice.
 	for _, path := range cmd.Paths {
 		// Parse path into bits.
 		logger.Printf("parsing: %s", path)
-		bits, err := cmd.parsePath(path)
-		if err != nil {
+		if err := cmd.importPath(path); err != nil {
 			return err
 		}
-
-		// Group bits by slice.
-		logger.Printf("grouping %d bits", len(bits))
-		bitsBySlice := pilosa.Bits(bits).GroupBySlice()
-		logger.Printf("grouped into %d slices", len(bitsBySlice))
-
-		// Parse path into bits.
-		for slice, bits := range bitsBySlice {
-			logger.Printf("importing slice: %d, n=%d", slice, len(bits))
-			if err := client.Import(cmd.Database, cmd.Frame, slice, bits); err != nil {
-				return err
-			}
-		}
 	}
 
 	return nil
 }
 
-// parsePath parses a path into bits.
-func (cmd *ImportCommand) parsePath(path string) ([]pilosa.Bit, error) {
-	var a []pilosa.Bit
+// importPath parses a path into bits and imports it to the server.
+func (cmd *ImportCommand) importPath(path string) error {
+	a := make([]pilosa.Bit, 0, cmd.BufferSize)
 
 	// Open file for reading.
 	f, err := os.Open(path)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	defer f.Close()
 
@@ -325,32 +320,64 @@ func (cmd *ImportCommand) parsePath(path string) ([]pilosa.Bit, error) {
 		if err == io.EOF {
 			break
 		} else if err != nil {
-			return nil, err
+			return err
 		}
 
 		// Ignore blank rows.
 		if record[0] == "" {
 			continue
 		} else if len(record) < 2 {
-			return nil, fmt.Errorf("bad column count on row %d: col=%d", rnum, len(record))
+			return fmt.Errorf("bad column count on row %d: col=%d", rnum, len(record))
 		}
 
 		// Parse bitmap id.
 		bitmapID, err := strconv.ParseUint(record[0], 10, 64)
 		if err != nil {
-			return nil, fmt.Errorf("invalid bitmap id on row %d: %q", rnum, record[0])
+			return fmt.Errorf("invalid bitmap id on row %d: %q", rnum, record[0])
 		}
 
 		// Parse bitmap id.
 		profileID, err := strconv.ParseUint(record[1], 10, 64)
 		if err != nil {
-			return nil, fmt.Errorf("invalid profile id on row %d: %q", rnum, record[1])
+			return fmt.Errorf("invalid profile id on row %d: %q", rnum, record[1])
 		}
 
 		a = append(a, pilosa.Bit{BitmapID: bitmapID, ProfileID: profileID})
+
+		// If we've reached the buffer size then import bits.
+		if len(a) == cmd.BufferSize {
+			if err := cmd.importBits(a); err != nil {
+				return err
+			}
+			a = a[:0]
+		}
 	}
 
-	return a, nil
+	// If there are still bits in the buffer then flush them.
+	if err := cmd.importBits(a); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// importPath parses a path into bits and imports it to the server.
+func (cmd *ImportCommand) importBits(bits []pilosa.Bit) error {
+	logger := log.New(cmd.Stderr, "", log.LstdFlags)
+
+	// Group bits by slice.
+	logger.Printf("grouping %d bits", len(bits))
+	bitsBySlice := pilosa.Bits(bits).GroupBySlice()
+
+	// Parse path into bits.
+	for slice, bits := range bitsBySlice {
+		logger.Printf("importing slice: %d, n=%d", slice, len(bits))
+		if err := cmd.Client.Import(cmd.Database, cmd.Frame, slice, bits); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
 
 // ExportCommand represents a command for bulk exporting data from a server.

executor.go

@@ -186,26 +186,41 @@ func (e *Executor) executeTopN(db string, c *pql.TopN, slices []uint64, opt *Exe
 }
 
 func (e *Executor) executeTopNSlices(db string, c *pql.TopN, slices []uint64, opt *ExecOptions) ([]Pair, error) {
-	var results []Pair
-	for node, nodeSlices := range e.slicesByNode(db, slices) {
-		// Execute locally if the hostname matches.
-		if node.Host == e.Host {
-			for _, slice := range nodeSlices {
-				pairs, err := e.executeTopNSlice(db, c, slice)
-				if err != nil {
-					return nil, err
-				}
-				results = Pairs(results).Add(pairs)
-			}
-			continue
-		}
+	slicesByNode := e.slicesByNode(db, slices)
 
-		// Otherwise execute remotely.
-		res, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nodeSlices, opt)
-		if err != nil {
-			return nil, err
+	type resp struct {
+		pairs []Pair
+		err   error
+	}
+	ch := make(chan resp, len(slicesByNode))
+
+	for node, nodeSlices := range slicesByNode {
+		go func(node *Node, nodeSlices []uint64) {
+			// Execute locally if the hostname matches.
+			if node.Host == e.Host {
+				pairs, err := e.executeTopNSlicesLocal(db, c, nodeSlices)
+				ch <- resp{pairs: pairs, err: err}
+				return
+			}
+
+			// Otherwise execute remotely.
+			res, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nodeSlices, opt)
+			if err != nil {
+				ch <- resp{err: err}
+				return
+			}
+			ch <- resp{pairs: res[0].([]Pair)}
+		}(node, nodeSlices)
+	}
+
+	// Collect results.
+	var results []Pair
+	for range slicesByNode {
+		r := <-ch
+		if r.err != nil {
+			return nil, r.err
 		}
-		results = Pairs(results).Add(res[0].([]Pair))
+		results = Pairs(results).Add(r.pairs)
 	}
 
 	// Sort final merged results.
@@ -219,6 +234,34 @@ func (e *Executor) executeTopNSlices(db string, c *pql.TopN, slices []uint64, op
 	return results, nil
 }
 
+func (e *Executor) executeTopNSlicesLocal(db string, c *pql.TopN, slices []uint64) ([]Pair, error) {
+	type resp struct {
+		pairs []Pair
+		err   error
+	}
+	ch := make(chan resp, len(slices))
+
+	// Execute TopN() in parallel across slices.
+	for _, slice := range slices {
+		go func(slice uint64) {
+			pairs, err := e.executeTopNSlice(db, c, slice)
+			ch <- resp{pairs: pairs, err: err}
+		}(slice)
+	}
+
+	// Collect results.
+	var results []Pair
+	for range slices {
+		r := <-ch
+		if r.err != nil {
+			return nil, r.err
+		}
+		results = Pairs(results).Add(r.pairs)
+	}
+
+	return results, nil
+}
+
 // executeTopNSlice executes a TopN call for a single slice.
 func (e *Executor) executeTopNSlice(db string, c *pql.TopN, slice uint64) ([]Pair, error) {
 	// Retrieve bitmap used to intersect.

fragment.go

@@ -27,7 +27,7 @@ import (
 
 const (
 	// SliceWidth is the number of profile IDs in a slice.
-	SliceWidth = 0x1000000 // 1048576
+	SliceWidth = 1048576
 
 	// SnapshotExt is the file extension used for an in-process snapshot.
 	SnapshotExt = ".snapshotting"

roaring/assembly_asm.go

@@ -3,27 +3,21 @@
 package roaring
 
 //go:noescape
-
 var useAsm = hasAsm()
 
 //go:noescape
-
 func popcntSliceAsm(s []uint64) uint64
 
 //go:noescape
-
 func popcntMaskSliceAsm(s, m []uint64) uint64
 
 //go:noescape
-
 func popcntAndSliceAsm(s, m []uint64) uint64
 
 //go:noescape
-
 func popcntOrSliceAsm(s, m []uint64) uint64
 
 //go:noescape
-
 func popcntXorSliceAsm(s, m []uint64) uint64
 
 //go:noescape

roaring/roaring.go

@@ -19,7 +19,7 @@ const (
 	headerSize = 4 + 4
 
 	// bitmapN is the number of values in a container.bitmap.
-	bitmapN = (1 << 24) / 64
+	bitmapN = (1 << 16) / 64
 )
 
 // Bitmap represents a roaring bitmap.
@@ -146,7 +146,7 @@ func (b *Bitmap) Max() uint64 {
 
 	hb := b.keys[len(b.keys)-1]
 	lb := b.containers[len(b.containers)-1].max()
-	return uint64(hb)<<24 | uint64(lb)
+	return uint64(hb)<<16 | uint64(lb)
 }
 
 // Count returns the number of bits set in the bitmap.
@@ -648,9 +648,9 @@ func (itr *Iterator) peek() uint64 {
 	key := itr.bitmap.keys[itr.i]
 	c := itr.bitmap.containers[itr.i]
 	if c.isArray() {
-		return uint64(key)<<24 | uint64(c.array[itr.j])
+		return uint64(key)<<16 | uint64(c.array[itr.j])
 	}
-	return uint64(key)<<24 | uint64(itr.j)
+	return uint64(key)<<16 | uint64(itr.j)
 }
 
 // BufIterator wraps an iterator to provide the ability to unread values.
@@ -704,7 +704,7 @@ func (itr *BufIterator) Unread() {
 }
 
 // The maximum size of array containers.
-const ArrayMaxSize = (1 << 20)
+const ArrayMaxSize = 4096
 
 // container represents a container for uint32 integers.
 //
@@ -1103,7 +1103,7 @@ func intersectionCountArrayBitmap(a, b *container) (n uint64) {
 }
 
 func intersectionCountBitmapBitmap(a, b *container) (n uint64) {
-	return popcntAndSliceGo(a.bitmap, b.bitmap)
+	return popcntAndSlice(a.bitmap, b.bitmap)
 }
 
 func intersect(a, b *container) *container {
@@ -1463,8 +1463,8 @@ func (op *op) UnmarshalBinary(data []byte) error {
 // size returns the encoded size of the op, in bytes.
 func (*op) size() int { return 1 + 8 + 4 }
 
-func highbits(v uint64) uint64 { return uint64(v >> 24) }
-func lowbits(v uint64) uint32  { return uint32(v & 0xFFFFFF) }
+func highbits(v uint64) uint64 { return uint64(v >> 16) }
+func lowbits(v uint64) uint32  { return uint32(v & 0xFFFF) }
 
 // search32 returns the index of v in a.
 func search32(a []uint32, value uint32) int {