Adam Drake

Faster command line tools with Go

Update

After some input from u/epiris on Reddit, I improved the code a bit further by changing the way bytestreams are scanned. Current fastest runtime is 0.308s for 10,512,769 rows, or about 34 Million rows per second. Since the file is 184 Megabytes, this is a processing speed of about 600 Megabytes per second, which is probably close to the read speed limit of my SSD.


It all started with a blog post, Faster Command Line Tools in D. It was at the top of Hacker News and got picked up by other sites as well. Then Euan Torano wrote Faster Command Line Tools in Nim. After reading about both, I found some comments in the Reddit post on the same topic, but in Go. Go is a great tool for many things, but I wasn’t sure what kind of performance I’d get out of this kind of tool.

TL;DR: I wrote a special-case version of the utility which processes the input in about 0.482 seconds, or about 22.8 million lines per second on my laptop, which is a processing rate of about 390 Megabytes per second.

The approximate performance of the other languages was as follows:

  • Python: 15 seconds
  • D (DMD): 2.4 seconds
  • D (LDC): 1.4 seconds
  • Nim: 1.4 seconds

  • My current-fastest Go version: 0.338s

Problem statement

The original version of the problem, as outlined in the post on D, was the following:

It’s a common programming task: Take a data file with fields separated by a delimiter (comma, tab, etc), and run a mathematical calculation involving several of the fields. Often these programs are one-time use scripts, other times they have longer shelf life. Speed is of course appreciated when the program is used more than a few times on large files.

The specific exercise we’ll explore starts with files having keys in one field, integer values in another. … Fields are delimited by a TAB, and there may be any number of fields on a line. The file name and field numbers of the key and value are passed as command line arguments.

Okay. We’ll have a delimited file (with tabs in this case), and we want to provide the filename, index for the key column, and index for the value column to our code. With those, our code should compute the sum of all the values for each specific key, and then output the key having the largest value. In SQL terms, this is a GROUP BY and a MAX. (Note that I did not do a version of this problem in SQL, but from what I’ve seen it is slower than the Go/Nim/D versions, though perhaps not the Python version.)

The data we will use is a file of n-grams from the Google Books dataset. The file is 184 Megabytes, uncompressed, and has a total of 10,512,769 lines.

First attempt

As a first try, we’ll use all the built-in Go libraries for string and file processing. We’ll keep the code general and high-level, and see what kind of speed we get.

Since I like to start with data structures, we’ll just use a simple map[string]int in this case, with the string keys being the keys from our file, and the integer values being the sum of all the values in the file for that particular key. In other words, it’s a for loop over the rows in the file, and an increment of a map value each time we go through the loop. Then at the end we’ll do a for loop over all the elements in the map and see which key has the largest value.

package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"math"
	"os"
	"strconv"
	"strings"
)

func processFile(filePath string, keyIndex, valueIndex int) {
	delim := "\t"

	fileHandle, err := os.Open(filePath)
	defer fileHandle.Close()

	maxFieldIndex := int(math.Max(float64(keyIndex), float64(valueIndex)))
	sumByKey := make(map[string]int)

	if err != nil {
		log.Fatal(err)
	}

	fileReader := bufio.NewScanner(fileHandle)

	for fileReader.Scan() {
		fields := strings.Split(strings.TrimSpace(fileReader.Text()), delim)

		if maxFieldIndex < len(fields) {
			value, err := strconv.Atoi(fields[valueIndex])
			if err != nil {
				log.Fatal(err)
			}
			sumByKey[fields[keyIndex]] += value
		}
	}

	maxValue := 0
	maxKey := ""
	for k, v := range sumByKey {
		if v > maxValue {
			maxValue = v
			maxKey = k
		}
	}
	fmt.Println("max_key:", maxKey, "sum:", sumByKey[maxKey])

}

func main() {
	filePath := flag.String("filePath", "", "Name of the file to parse")
	keyIndex := flag.Int("keyIndex", 0, "Index of key (0 is first position)")
	valueIndex := flag.Int("valueIndex", 0, "Index of value (0 is first position)")

	flag.Parse()
	processFile(*filePath, *keyIndex, *valueIndex)

}

The processFile() function does all the work, and that makes it easy for us to write a benchmark.

package main

import (
	"testing"
)

func Benchmark_processFile(b *testing.B) {
	for i := 0; i < b.N; i++ {
		processFile("../ngrams.tsv", 1, 2)
	}
}

Results

The first version works, but it’s not very fast. We can see its performance by running go test -cpuprofile cpu.prof -memprofile mem.prof -bench . both to get a benchmark, and to see which parts are slow.

max_key: 2006 sum: 22569013
Benchmark_processFile-4   	       1	3486100959 ns/op
PASS
ok  	github.com/adamdrake/faster-command-line-tools-in-nim/Go/v1	3.491s

Our first version takes about 3.49 seconds to run, which isn’t bad, but isn’t great. Why so slow? Let’s go tool pprof cpu.prof and find out.

Entering interactive mode (type "help" for commands)
(pprof) text
2250ms of 3490ms total (64.47%)
Dropped 33 nodes (cum <= 17.45ms)
Showing top 10 nodes out of 65 (cum >= 160ms)
      flat  flat%   sum%        cum   cum%
     600ms 17.19% 17.19%     1530ms 43.84%  strings.genSplit
     330ms  9.46% 26.65%      790ms 22.64%  runtime.mallocgc
     260ms  7.45% 34.10%      260ms  7.45%  runtime.heapBitsSetType
     200ms  5.73% 39.83%      200ms  5.73%  runtime.indexbytebody
     200ms  5.73% 45.56%      310ms  8.88%  runtime.mapaccess1_faststr
     170ms  4.87% 50.43%      410ms 11.75%  runtime.mapassign
     160ms  4.58% 55.01%      160ms  4.58%  runtime.aeshashbody
     120ms  3.44% 58.45%      400ms 11.46%  strings.Count
     110ms  3.15% 61.60%     3440ms 98.57%  github.com/adamdrake/faster-command-line-tools-in-nim/Go/v1.processFile
     100ms  2.87% 64.47%      160ms  4.58%  strconv.ParseInt

It seems like we spend a lot of time on genSplit(), which is the General Split command invoked by strings.Split() in our code. To see this, we can type list processFile in pprof to get the line-by-line timing for that function. Here’s a selection of the output of list processFile.

         .      240ms     29:	for fileReader.Scan() {
         .      2.19s     30:		fields := strings.Split(strings.TrimSpace(fileReader.Text()), delim)
         .          .     31:
         .          .     32:		if maxFieldIndex < len(fields) {
      10ms      190ms     33:			value, err := strconv.Atoi(fields[valueIndex])
         .          .     34:			if err != nil {
         .          .     35:				log.Fatal(err)
         .          .     36:			}
      80ms      800ms     37:			sumByKey[fields[keyIndex]] += value
         .          .     38:		}
         .          .     39:	}

We see that out of the 3.49s we’re spending processing our file, 2.19s of that is spent just splitting strings. To confirm that genSplit() is in strings.Split() we can do the same thing with list Split and get all the details on where the time is going.

Let’s assume that we aren’t going to dive in and start making changes to the standard library. What should we do to improve the speed? The natural thing for most Go programmers is to turn to channels and goroutines. Spoiler alert: that is slower.

V2: Channels and Goroutines

Channels and goroutines are great, but often overused. They have a startup cost, and many people are not aware of the fact that channels are a data structure that includes a mutex, and therefore subject to lock contention (i.e., they can be slow). Regardless, let’s see how the performance fares if we go that route. We’ll keep the map and other data in a struct, which we’ll protect with a sync.Mutex so that we have thread-safe write access to the map (recall: in Go, a map is NOT thread-safe).

package main

import (
	"flag"
	"fmt"
	"log"
	"math"
	"strconv"
	"strings"
	"sync"

	fstream "github.com/adamdrake/gofstream"
)

type kv struct {
	keyIndex, valueIndex int
	store                map[string]int
	sync.Mutex
}

func worker(rows chan string, data *kv, wg *sync.WaitGroup) {
	defer wg.Done()
	keyIndex := data.keyIndex
	valueIndex := data.valueIndex
	maxFieldIndex := int(math.Max(float64(keyIndex), float64(valueIndex)))
	sumByKey := make(map[string]int)
	for r := range rows {
		fields := strings.Split(strings.TrimSpace(r), "\t")

		if maxFieldIndex < len(fields) {
			value, err := strconv.Atoi(fields[data.valueIndex])
			if err != nil {
				log.Fatal(err)
			}
			sumByKey[fields[keyIndex]] += value
		}
	}
	data.Lock()
	for k, v := range sumByKey {
		data.store[k] += v

	}
	data.Unlock()
}

func processFile(filePath string, keyIndex, valueIndex int) {

	rows, err := fstream.New(filePath, 100000)
	if err != nil {
		log.Fatal(err)
	}

	data := &kv{keyIndex: keyIndex, valueIndex: valueIndex, store: make(map[string]int)}
	wg := &sync.WaitGroup{}
	for i := 0; i < 4; i++ {
		wg.Add(1)
		go worker(rows, data, wg)
	}
	wg.Wait()

	maxValue := 0
	maxKey := ""
	for k, v := range data.store {
		if v > maxValue {
			maxValue = v
			maxKey = k
		}
	}
	fmt.Println("max_key:", maxKey, "sum:", data.store[maxKey])

}

func main() {
	filePath := flag.String("filePath", "", "Name of the file to parse")
	keyIndex := flag.Int("keyIndex", 0, "Index of key (0 is first position)")
	valueIndex := flag.Int("valueIndex", 0, "Index of value (0 is first position)")

	flag.Parse()
	processFile(*filePath, *keyIndex, *valueIndex)

}

This code has the same processFile() function as before, but this time it creates a streaming reader for the file, and streams each row of the file over a channel. We then give the channel to a worker thread, and the workers do the same thing we did before. We also have some additional complexity in the form of a sync.WaitGroup to make sure we actually wait and allow the workers to do their job before main() returns. What’s the performance? Not great. In fact, it’s worse, at about 4.5s.

max_key: 2006 sum: 22569013
Benchmark_processFile-4   	       1	4476065569 ns/op
PASS
ok  	github.com/adamdrake/faster-command-line-tools-in-nim/Go/v4	4.482s

I won’t go into the pprof output for this one, but you can see that the reason it’s slower is because of runtime.procyield, due to all our goroutines. I probably could have tried to do some things with sync/atomic instead of using a mutex on the struct, but spinning up goroutines for such a small file doesn’t make much sense in this case anyway. If the file was larger, perhaps we’d get the benefit of using all cores, but for input of this size it just isn’t worth it.

V4: Stop using strings

Since it seems that a lot of the overhead is in parsing and splitting strings, why not just operate on the underlying bytes instead? Credit to valyala for their version, which reads and splits on bytes instead of string. Here’s my version of their code, with an additional twist from a Stack Overflow post which allows us to use a faster string to integer parsing method, assuming that the input values we are summing are always positive integers.

package main

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"os"
	"unsafe"
)

func processFile(file *os.File, keyField, valueField int) {
	var sumByKey = make(map[string]int)
	maxField := keyField
	if valueField > maxField {
		maxField = valueField
	}

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		line := scanner.Bytes()
		key, val := getKeyVal(line, keyField, valueField, maxField)
		sumByKey[key] += val
	}
	var k string
	v := -1
	for key, val := range sumByKey {
		if val > v {
			k = key
			v = val
		}
	}
	fmt.Printf("max_key: %s sum: %d", k, v)

}

func main() {

	file, _ := os.Open(os.Args[1])
	defer file.Close()

	keyField, _ := atoi([]byte(os.Args[2]))
	valueField, _ := atoi([]byte(os.Args[3]))
	processFile(file, keyField, valueField)

}

func getKeyVal(line []byte, keyField, valueField, maxField int) (string, int) {
	var i int
	var tabIndex int
	var k []byte
	var v []byte
	var field []byte

	for i <= maxField && len(line) > 0 {
		tabIndex = bytes.IndexByte(line, '\t') // returns -1 if not found
		if tabIndex < 0 {
			field = line
			line = nil
		} else {
			field = line[:tabIndex]
			line = line[tabIndex+1:]
		}
		switch i {
		case keyField:
			k = field
		case valueField:
			v = field
		}
		i++
	}
	val, _ := atoi(v)
	return string(k), val
}

var errAtoi = errors.New("invalid number")

func atoi(input []byte) (int, error) {
	val := 0
	for i := 0; i < len(input); i++ {
		char := input[i]
		if char < '0' || char > '9' {
			return 0, errAtoi
		}
		val = val*10 + int(char) - '0'
	}
	return val, nil
}

Now, instead of reading in the line as a string and then splitting the string on the delimiter with strings.Split(), we’re reading in the line as a bytestring, and using the bytes.IndexByte() function to find the index of the delimiter in the bytestring. We loop through the bytestring by finding the current delimiter (which corresponds to the command line arguments we received for key index and value index), and sum the values as always. How much do we save by operating on the bytes instead of the strings? A lot.

max_key: 2006 sum: 22569013
Benchmark_processFile-4   	       1	1517321373 ns/op
PASS

Great! This approach cut our runtime by about 60%, and it’s still pretty usable as production code. It’s not too customized, it still fits the requirements of the original problem definition of key column and value column being command line arguments, and it’s pretty readable. Where is the time mostly spent?

(pprof) text
1200ms of 1500ms total (80.00%)
Showing top 10 nodes out of 41 (cum >= 50ms)
      flat  flat%   sum%        cum   cum%
     260ms 17.33% 17.33%      670ms 44.67%  github.com/adamdrake/faster-command-line-tools-in-nim/Go/v5.getKeyVal
     200ms 13.33% 30.67%      330ms 22.00%  runtime.mapaccess1_faststr
     170ms 11.33% 42.00%      170ms 11.33%  runtime.indexbytebody
     120ms  8.00% 50.00%      260ms 17.33%  runtime.mapassign
     100ms  6.67% 56.67%      100ms  6.67%  runtime.aeshashbody
      90ms  6.00% 62.67%      120ms  8.00%  runtime.mallocgc
      70ms  4.67% 67.33%     1500ms   100%  github.com/adamdrake/faster-command-line-tools-in-nim/Go/v5.processFile
      70ms  4.67% 72.00%       70ms  4.67%  runtime.memeqbody
      70ms  4.67% 76.67%       80ms  5.33%  syscall.Syscall
      50ms  3.33% 80.00%       50ms  3.33%  github.com/adamdrake/faster-command-line-tools-in-nim/Go/v5.atoi

Most of the time is spent splitting the bytes and getting the key/value pairs, as expected. Let’s list getKeyVal to take a closer look.

260ms      670ms (flat, cum) 44.67% of Total
         .          .     51:	var tabIndex int
         .          .     52:	var k []byte
         .          .     53:	var v []byte
         .          .     54:	var field []byte
         .          .     55:
      20ms       20ms     56:	for i <= maxField && len(line) > 0 {
      20ms      160ms     57:		tabIndex = bytes.IndexByte(line, '\t') // returns -1 if not found
      20ms       20ms     58:		if tabIndex < 0 {
         .          .     59:			field = line
         .          .     60:			line = nil
         .          .     61:		} else {
      10ms       10ms     62:			field = line[:tabIndex]
      60ms       60ms     63:			line = line[tabIndex+1:]
         .          .     64:		}
         .          .     65:		switch i {
      10ms       10ms     66:		case keyField:
         .          .     67:			k = field
      10ms       10ms     68:		case valueField:
         .          .     69:			v = field
         .          .     70:		}
      10ms       10ms     71:		i++
         .          .     72:	}
      40ms       90ms     73:	val, _ := atoi(v)
      60ms      280ms     74:	return string(k), val
         .          .     75:}

We can see that the most time is spent getting the index of the bytes, and some on string casting for the key (so that we can use it as the map key). In a practical scenario, we’d probably stop here since, as mentioned above, this is probably about as fast as the code will get while still being sufficiently general, readable, etc. We can do better though, if we break the requirements of the original comparison.

V5: Let’s go nuts

Let’s put some of the practical concerns aside for the moment and focus on what we know about the problem. First, we are trusting our input data, which we’d normally be more careful about in a real-world situation. With that in mind, we can change the definitions of our atoi() function so that it doesn’t do any error checking, and so that it only operates on []byte and returns int. Here’s what the simplified atoi looks like.

func atoi(s []byte) int {
	i := 0
	x := 0
	for ; i < len(s); i++ {
		c := s[i]
		x = x*10 + int(c) - '0'
	}
	return x
}

We also know that the keys and values are always integers, so we could use an int type instead of string in the map keys, and we could also remove the string(k) cast in the getKeyVal() implementation above. However, we know in this case that we want the value between the first and second separator, and between the second and third separator, every time. So in this contrived instance, we don’t actually need to iterate and check delimiters. We can just treat the entire line like the bytestring it is, and only find the indexes for the first, second, and third instances of the delimiter.

Another important fact is that, for this dataset, we know the largest key. Therefore, instead of a map that we previously saw was costing us about 800ms to increment, we can use an array. In this case, our integer key will be the index, and the value at that index will be the sum for that key. With that in mind, we have a highly customised, but SUPER FAST version below.

func main() {

	file, err := os.Open(os.Args[1])
	defer file.Close()
	if err != nil {
		log.Fatal(err)
	}
	processFile(file)
}

func processFile(file *os.File) {
	var sumByKey [2009]int

	scanner := bufio.NewScanner(file)
	var key int
	var val int

	for scanner.Scan() {
		line := scanner.Bytes()
		firstTab := bytes.IndexByte(line, '\t')
		secondTab := bytes.IndexByte(line[firstTab+1:], '\t') + firstTab + 1
		thirdTab := bytes.IndexByte(line[secondTab+1:], '\t') + secondTab + 1
		key = atoi(line[firstTab+1 : secondTab])
		val = atoi(line[secondTab+1 : thirdTab])
		sumByKey[key] += val
	}
	var k int
	var v int
	for i, val := range sumByKey {
		if val > v {
			k = i
			v = val
		}
	}
	fmt.Printf("max_key: %d sum: %d\n", k, v)
}

What’s the pprof result?

Benchmark_processFile-4   	max_key: 2006 sum: 22569013
max_key: 2006 sum: 22569013
max_key: 2006 sum: 22569013
max_key: 2006 sum: 22569013
max_key: 2006 sum: 22569013
       3	 482369483 ns/op

It’s crazy fast! This version finishes in 0.482s (482ms) on my machine. Where is the most time spent now? (Note that because the benchmark ran multiple times in order to get a stable timing for the run, the pprof results are the sum of all the runs.)

      20ms      1.26s     38:	for scanner.Scan() {
      70ms       70ms     39:		line := scanner.Bytes()
      40ms      200ms     40:		firstTab := bytes.IndexByte(line, '\t')
     190ms      330ms     41:		secondTab := bytes.IndexByte(line[firstTab+1:], '\t') + firstTab + 1
     260ms      610ms     42:		thirdTab := bytes.IndexByte(line[secondTab+1:], '\t') + secondTab + 1
      80ms      160ms     43:		key = atoi(line[firstTab+1 : secondTab])
      40ms      180ms     44:		val = atoi(line[secondTab+1 : thirdTab])
      70ms       70ms     45:		sumByKey[key] += val
         .          .     46:	}

Most of the time is spent indexing into the bytestring, which is what we expected. There may be some clever ways to reduce this, in addition to some data structures which would allow us to do the incrementing faster, but at this point we’re pretty far down the path of optimizing past production usage, and hand-tuned Fortran/assembly isn’t really the goal. In the end though, we processed a 184MB file, with 10,512,769 rows, in 0.482s.

Happy optimizing!

Updated code:

Since the article was posted on Reddit and Hackernews, there was a user (Epiris) on Reddit who suggested some improvements. After their suggestions, below is the code and benchmarks for the updated version. It’s important to note that a personal requirement I have for the code is that it cannot read all data into memory, since we must assume that the data stream could be arbitrarily large.

func processLine(b []byte) (int, int) {
	key := 0
	val := 0
	i := 0
	for b[i] != '\t' {
		i++
	}
	for i++; b[i] != '\t'; i++ {
		key = key*10 + int(b[i]) - '0'
	}
	for i++; b[i] != '\t'; i++ {
		val = val*10 + int(b[i]) - '0'
	}
	return key, val
}

func processFile(file *os.File) (int, int) {
	var sumByKey [2009]int

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Bytes()
		k1, v1 := processLine(line)
		sumByKey[k1] += v1
	}
	var k int
	var v int
	for i, val := range sumByKey {
		if val > v {
			k = i
			v = val
		}
	}
	return k, v

}

And here is the associated benchmark.

func Benchmark_processFile(b *testing.B) {
	file, err := os.Open("../ngrams.tsv")
	defer file.Close()
	if err != nil {
		b.Fatal(err)
	}

	data, err := ioutil.ReadAll(file)
	if err != nil {
		b.Fatal(err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		k, v := processFile(data)
		if k != 2006 || v != 22569013 {
			b.Fatalf(`bad result %v | %v`, k, v)
		}
	}
}

The output of the benchmarking:

Benchmark_processFile-4   	       5	 309631168 ns/op	    4152 B/op	       4 allocs/op