Repository: segmentio/parquet-go
Pull request #273 (2022-07-19 20:01:00)

Improve BYTE_ARRAY dictionaries
This PR contributes to #226 by optimizing the `byteArrayDictionary` type.

I made multiple attempts to improve throughput of probing operations for string values (reflected in this branch history) but without being able to significantly increase throughput in the way that I did for probing fixed-length values (e.g. int32). The increased complexity did not justify the returns, so we will stick with `map[string]int32` in this case. 

Overall, measurable but minor improvements, bug fixes, and simplification of the code.

```
name                                    old time/op  new time/op  delta
Dictionary/Insert/BYTE_ARRAY/N=100      1.40µs ± 0%  1.36µs ± 1%  -2.88%  (p=0.000 n=8+9)
Dictionary/Insert/BYTE_ARRAY/N=1000     13.8µs ± 1%  13.6µs ± 1%  -1.73%  (p=0.000 n=10+9)
Dictionary/Insert/BYTE_ARRAY/N=10000     414µs ± 2%   403µs ± 0%  -2.72%  (p=0.000 n=10+10)
Dictionary/Insert/BYTE_ARRAY/N=100000   5.58ms ± 4%  5.57ms ± 2%    ~     (p=0.579 n=10+10)
Dictionary/Insert/BYTE_ARRAY/N=1000000   105ms ± 2%   105ms ± 2%    ~     (p=0.436 n=10+10)

name                                    old value/s  new value/s  delta
Dictionary/Insert/BYTE_ARRAY/N=100       71.3M ± 1%   73.5M ± 1%  +3.04%  (p=0.000 n=9+9)
Dictionary/Insert/BYTE_ARRAY/N=1000      72.3M ± 1%   73.6M ± 1%  +1.77%  (p=0.000 n=10+9)
Dictionary/Insert/BYTE_ARRAY/N=10000     24.1M ± 2%   24.8M ± 0%  +2.79%  (p=0.000 n=10+10)
Dictionary/Insert/BYTE_ARRAY/N=100000    17.9M ± 4%   18.0M ± 2%    ~     (p=0.579 n=10+10)
Dictionary/Insert/BYTE_ARRAY/N=1000000   9.56M ± 2%   9.51M ± 2%    ~     (p=0.436 n=10+10)
```

achille-roussel: Thanks for the review!

Patch:
diff --git a/dictionary.go b/dictionary.go
index d800599..7021613 100644
--- a/dictionary.go
+++ b/dictionary.go
@@ -3,6 +3,7 @@ package parquet
 import (
 	"io"
 	"math/bits"
+	"strings"
 	"unsafe"
 
 	"github.com/segmentio/parquet-go/deprecated"
@@ -159,10 +160,11 @@ func (d *booleanDictionary) insert(indexes []int32, rows sparse.Array) {
 		d.bits = plain.AppendBoolean(d.bits, int(d.table[1]), true)
 	}
 
+	values := rows.Uint8Array()
 	dict := d.table
 
 	for i := 0; i < rows.Len(); i++ {
-		v := *(*byte)(rows.Index(i)) & 1
+		v := values.Index(i) & 1
 		indexes[i] = dict[v]
 	}
 }
@@ -242,7 +244,7 @@ func (d *int32Dictionary) Insert(indexes []int32, values []Value) {
 }
 
 func (d *int32Dictionary) init(indexes []int32) {
-	d.table = hashprobe.NewInt32Table(cap(d.values), hashprobeTableMaxLoad)
+	d.table = hashprobe.NewInt32Table(len(d.values), hashprobeTableMaxLoad)
 
 	n := min(len(d.values), len(indexes))
 
@@ -342,7 +344,7 @@ func (d *int64Dictionary) Insert(indexes []int32, values []Value) {
 }
 
 func (d *int64Dictionary) init(indexes []int32) {
-	d.table = hashprobe.NewInt64Table(cap(d.values), hashprobeTableMaxLoad)
+	d.table = hashprobe.NewInt64Table(len(d.values), hashprobeTableMaxLoad)
 
 	n := min(len(d.values), len(indexes))
 
@@ -439,7 +441,7 @@ func (d *int96Dictionary) insertValues(indexes []int32, count int, valueAt func(
 	_ = indexes[:count]
 
 	if d.hashmap == nil {
-		d.hashmap = make(map[deprecated.Int96]int32, cap(d.values))
+		d.hashmap = make(map[deprecated.Int96]int32, len(d.values))
 		for i, v := range d.values {
 			d.hashmap[v] = int32(i)
 		}
@@ -524,7 +526,7 @@ func (d *floatDictionary) Insert(indexes []int32, values []Value) {
 }
 
 func (d *floatDictionary) init(indexes []int32) {
-	d.table = hashprobe.NewFloat32Table(cap(d.values), hashprobeTableMaxLoad)
+	d.table = hashprobe.NewFloat32Table(len(d.values), hashprobeTableMaxLoad)
 
 	n := min(len(d.values), len(indexes))
 
@@ -611,7 +613,7 @@ func (d *doubleDictionary) Insert(indexes []int32, values []Value) {
 }
 
 func (d *doubleDictionary) init(indexes []int32) {
-	d.table = hashprobe.NewFloat64Table(cap(d.values), hashprobeTableMaxLoad)
+	d.table = hashprobe.NewFloat64Table(len(d.values), hashprobeTableMaxLoad)
 
 	n := min(len(d.values), len(indexes))
 
@@ -671,28 +673,28 @@ func (d *doubleDictionary) Page() BufferedPage {
 
 type byteArrayDictionary struct {
 	byteArrayPage
+	table   map[string]int32
 	offsets []uint32
-	hashmap map[string]int32
 }
 
 func newByteArrayDictionary(typ Type, columnIndex int16, numValues int32, values []byte) *byteArrayDictionary {
 	d := &byteArrayDictionary{
-		offsets: make([]uint32, 0, numValues),
 		byteArrayPage: byteArrayPage{
 			typ:         typ,
 			values:      values,
-			numValues:   numValues,
 			columnIndex: ^columnIndex,
 		},
+		offsets: make([]uint32, 0, numValues),
 	}
 
 	for i := 0; i < len(values); {
-		n := plain.ByteArrayLength(values[i:])
 		d.offsets = append(d.offsets, uint32(i))
+		n := plain.ByteArrayLength(values[i:])
 		i += plain.ByteArrayLengthSize
 		i += n
 	}
 
+	d.numValues = int32(len(d.offsets))
 	return d
 }
 
@@ -709,36 +711,46 @@ func (d *byteArrayDictionary) Insert(indexes []int32, values []Value) {
 	d.insert(indexes, makeArrayValue(values, unsafe.Offsetof(model.ptr)))
 }
 
-func (d *byteArrayDictionary) insert(indexes []int32, rows sparse.Array) {
-	_ = indexes[:rows.Len()]
+func (d *byteArrayDictionary) init() {
+	d.table = make(map[string]int32, d.numValues)
 
-	if d.hashmap == nil {
-		d.hashmap = make(map[string]int32, cap(d.offsets))
-		for index, offset := range d.offsets {
-			d.hashmap[string(d.valueAt(offset))] = int32(index)
-		}
+	for i := 0; i < len(d.values); {
+		n := plain.ByteArrayLength(d.values[i:])
+		i += plain.ByteArrayLengthSize
+		s := d.values[i : i+n]
+		i += n
+		d.table[string(s)] = int32(len(d.table))
 	}
+}
 
-	for i := 0; i < rows.Len(); i++ {
-		value := *(*string)(rows.Index(i))
+func (d *byteArrayDictionary) insert(indexes []int32, rows sparse.Array) {
+	if d.table == nil {
+		d.init()
+	}
 
-		index, exists := d.hashmap[value]
+	values := rows.StringArray()
+
+	for i := range indexes {
+		value := values.Index(i)
+
+		index, exists := d.table[value]
 		if !exists {
-			index = int32(len(d.offsets))
-			value = d.append(value)
-			d.hashmap[value] = index
+			value = cloneString(value)
+			index = d.numValues
+			d.numValues++
+			d.table[value] = index
+			d.offsets = append(d.offsets, uint32(len(d.values)))
+			d.values = plain.AppendByteArrayString(d.values, value)
 		}
 
 		indexes[i] = index
 	}
 }
 
-func (d *byteArrayDictionary) append(value string) string {
-	offset := len(d.values)
-	d.values = plain.AppendByteArrayString(d.values, value)
-	d.offsets = append(d.offsets, uint32(offset))
-	d.numValues++
-	return string(d.values[offset+plain.ByteArrayLengthSize : len(d.values)])
+func cloneString(s string) string {
+	b := new(strings.Builder)
+	b.WriteString(s)
+	return b.String()
 }
 
 func (d *byteArrayDictionary) Lookup(indexes []int32, values []Value) {
@@ -782,7 +794,10 @@ func (d *byteArrayDictionary) Reset() {
 	d.offsets = d.offsets[:0]
 	d.values = d.values[:0]
 	d.numValues = 0
-	d.hashmap = nil
+
+	for k := range d.table {
+		delete(d.table, k)
+	}
 }
 
 func (d *byteArrayDictionary) Page() BufferedPage {
@@ -933,7 +948,7 @@ func (d *uint32Dictionary) Insert(indexes []int32, values []Value) {
 }
 
 func (d *uint32Dictionary) init(indexes []int32) {
-	d.table = hashprobe.NewUint32Table(cap(d.values), hashprobeTableMaxLoad)
+	d.table = hashprobe.NewUint32Table(len(d.values), hashprobeTableMaxLoad)
 
 	n := min(len(d.values), len(indexes))
 
@@ -1020,7 +1035,7 @@ func (d *uint64Dictionary) Insert(indexes []int32, values []Value) {
 }
 
 func (d *uint64Dictionary) init(indexes []int32) {
-	d.table = hashprobe.NewUint64Table(cap(d.values), hashprobeTableMaxLoad)
+	d.table = hashprobe.NewUint64Table(len(d.values), hashprobeTableMaxLoad)
 
 	n := min(len(d.values), len(indexes))
 
@@ -1138,7 +1153,7 @@ func (d *be128Dictionary) Insert(indexes []int32, values []Value) {
 }
 
 func (d *be128Dictionary) init(indexes []int32) {
-	d.table = hashprobe.NewUint128Table(cap(d.values), 0.75)
+	d.table = hashprobe.NewUint128Table(len(d.values), 0.75)
 
 	n := min(len(d.values), len(indexes))
 
diff --git a/hashprobe/hashprobe.go b/hashprobe/hashprobe.go
index 9e61aab..9e4a597 100644
--- a/hashprobe/hashprobe.go
+++ b/hashprobe/hashprobe.go
@@ -68,6 +68,13 @@ func init() {
 	prngSource = rand.NewSource(seed).(rand.Source64)
 }
 
+func tableSizeAndMaxLen(groupSize, numValues int, maxLoad float64) (size, maxLen int) {
+	n := int(math.Ceil((1 / maxLoad) * float64(numValues)))
+	size = nextPowerOf2((n + (groupSize - 1)) / groupSize)
+	maxLen = int(math.Ceil(maxLoad * float64(groupSize*size)))
+	return
+}
+
 func nextPowerOf2(n int) int {
 	return 1 << (64 - bits.LeadingZeros64(uint64(n-1)))
 }
@@ -185,13 +192,12 @@ func (t *table32) size() int {
 }
 
 func (t *table32) init(cap int, maxLoad float64) {
-	m := int(math.Ceil((1 / maxLoad) * float64(cap)))
-	n := nextPowerOf2((m + (table32GroupSize - 1)) / table32GroupSize)
+	size, maxLen := tableSizeAndMaxLen(table32GroupSize, cap, maxLoad)
 	*t = table32{
-		maxLen:  int(math.Ceil(maxLoad * float64(table32GroupSize*n))),
+		maxLen:  maxLen,
 		maxLoad: maxLoad,
 		seed:    randSeed(),
-		table:   make([]table32Group, n),
+		table:   make([]table32Group, size),
 	}
 }
 
@@ -431,13 +437,12 @@ func (t *table64) size() int {
 }
 
 func (t *table64) init(cap int, maxLoad float64) {
-	m := int(math.Ceil((1 / maxLoad) * float64(cap)))
-	n := nextPowerOf2((m + (table64GroupSize - 1)) / table64GroupSize)
+	size, maxLen := tableSizeAndMaxLen(table64GroupSize, cap, maxLoad)
 	*t = table64{
-		maxLen:  int(math.Ceil(maxLoad * float64(table64GroupSize*n))),
+		maxLen:  maxLen,
 		maxLoad: maxLoad,
 		seed:    randSeed(),
-		table:   make([]table64Group, n),
+		table:   make([]table64Group, size),
 	}
 }
 
@@ -622,14 +627,13 @@ func makeTable128(cap int, maxLoad float64) (t table128) {
 }
 
 func (t *table128) init(cap int, maxLoad float64) {
-	m := int(math.Ceil((1 / maxLoad) * float64(cap)))
-	n := nextPowerOf2(m)
+	size, maxLen := tableSizeAndMaxLen(1, cap, maxLoad)
 	*t = table128{
-		cap:     n,
-		maxLen:  int(math.Ceil(maxLoad * float64(n))),
+		cap:     size,
+		maxLen:  maxLen,
 		maxLoad: maxLoad,
 		seed:    randSeed(),
-		table:   make([]byte, 16*n+4*n),
+		table:   make([]byte, 16*size+4*size),
 	}
 }
 
diff --git a/hashprobe/hashprobe_purego.go b/hashprobe/hashprobe_purego.go
index 1c81781..13e0530 100644
--- a/hashprobe/hashprobe_purego.go
+++ b/hashprobe/hashprobe_purego.go
@@ -2,7 +2,9 @@
 
 package hashprobe
 
-import "github.com/segmentio/parquet-go/sparse"
+import (
+	"github.com/segmentio/parquet-go/sparse"
+)
 
 func multiProbe32(table []table32Group, numKeys int, hashes []uintptr, keys sparse.Uint32Array, values []int32) int {
 	return multiProbe32Default(table, numKeys, hashes, keys, values)
diff --git a/page.go b/page.go
index 9405144..4462116 100644
--- a/page.go
+++ b/page.go
@@ -1002,9 +1002,9 @@ func (page *byteArrayPage) Buffer() BufferedPage { return page }
 
 func (page *byteArrayPage) valueAt(offset uint32) []byte {
 	length := binary.LittleEndian.Uint32(page.values[offset:])
-	j := 4 + offset
-	k := 4 + offset + length
-	return page.values[j:k:k]
+	i := 4 + offset
+	j := 4 + offset + length
+	return page.values[i:j:j]
 }
 
 func (page *byteArrayPage) min() (min []byte) {
diff --git a/sparse/array.go b/sparse/array.go
index 5e4208b..401837a 100644
--- a/sparse/array.go
+++ b/sparse/array.go
@@ -24,6 +24,7 @@ func (a Array) Uint16Array() Uint16Array   { return Uint16Array{a.array} }
 func (a Array) Uint32Array() Uint32Array   { return Uint32Array{a.array} }
 func (a Array) Uint64Array() Uint64Array   { return Uint64Array{a.array} }
 func (a Array) Uint128Array() Uint128Array { return Uint128Array{a.array} }
+func (a Array) StringArray() StringArray   { return StringArray{a.array} }
 
 type array struct {
 	ptr unsafe.Pointer
@@ -270,3 +271,19 @@ func (a Uint128Array) Uint16Array() Uint16Array    { return Uint16Array{a.array}
 func (a Uint128Array) Uint32Array() Uint32Array    { return Uint32Array{a.array} }
 func (a Uint128Array) Uint64Array() Uint64Array    { return Uint64Array{a.array} }
 func (a Uint128Array) UnsafeArray() Array          { return Array{a.array} }
+
+type StringArray struct{ array }
+
+func MakeStringArray(values []string) StringArray {
+	const sizeOfString = unsafe.Sizeof("")
+	return StringArray{makeArray(*(*unsafe.Pointer)(unsafe.Pointer(&values)), uintptr(len(values)), sizeOfString)}
+}
+
+func UnsafeStringArray(base unsafe.Pointer, length int, offset uintptr) StringArray {
+	return StringArray{makeArray(base, uintptr(length), offset)}
+}
+
+func (a StringArray) Len() int                   { return int(a.len) }
+func (a StringArray) Index(i int) string         { return *(*string)(a.index(i)) }
+func (a StringArray) Slice(i, j int) StringArray { return StringArray{a.slice(i, j)} }
+func (a StringArray) UnsafeArray() Array         { return Array{a.array} }
diff --git a/sparse/gather.go b/sparse/gather.go
index ac83897..05d4500 100644
--- a/sparse/gather.go
+++ b/sparse/gather.go
@@ -26,6 +26,8 @@ func GatherUint64(dst []uint64, src Uint64Array) int { return gather64(dst, src)
 
 func GatherUint128(dst [][16]byte, src Uint128Array) int { return gather128(dst, src) }
 
+func GatherString(dst []string, src StringArray) int { return gatherString(dst, src) }
+
 func min(a, b int) int {
 	if a < b {
 		return a
diff --git a/sparse/gather_amd64.go b/sparse/gather_amd64.go
index 8b24cee..a000966 100644
--- a/sparse/gather_amd64.go
+++ b/sparse/gather_amd64.go
@@ -2,7 +2,11 @@
 
 package sparse
 
-import "golang.org/x/sys/cpu"
+import (
+	"unsafe"
+
+	"golang.org/x/sys/cpu"
+)
 
 func gatherBits(dst []byte, src Uint8Array) int {
 	n := min(len(dst)*8, src.Len())
@@ -81,3 +85,11 @@ func gather64AVX2(dst []uint64, src Uint64Array)
 
 //go:noescape
 func gather128(dst [][16]byte, src Uint128Array) int
+
+// On AMD64 the `string` type is made of a 64 bit pointer and a 64 bit integer,
+// for a total of 128 bit which means that we can use the same algorithm as
+// gather128 to lookup string values, since the memory layout of the arguments is
+// the same.
+func gatherString(dst []string, src StringArray) int {
+	return gather128(*(*[][16]byte)(unsafe.Pointer(&dst)), *(*Uint128Array)(unsafe.Pointer(&src)))
+}
diff --git a/sparse/gather_purego.go b/sparse/gather_purego.go
index 2f25c44..65f578e 100644
--- a/sparse/gather_purego.go
+++ b/sparse/gather_purego.go
@@ -70,3 +70,13 @@ func gather128(dst [][16]byte, src Uint128Array) int {
 
 	return n
 }
+
+func gatherString(dst []string, src StringArray) int {
+	n := min(len(dst), src.Len())
+
+	for i := range dst[:n] {
+		dst[i] = src.Index(i)
+	}
+
+	return n
+}


Test patch:
diff --git a/dictionary_test.go b/dictionary_test.go
index 4ab350b..b2ffb26 100644
--- a/dictionary_test.go
+++ b/dictionary_test.go
@@ -25,7 +25,7 @@ var dictionaryTypes = [...]parquet.Type{
 
 func TestDictionary(t *testing.T) {
 	for _, typ := range dictionaryTypes {
-		for _, numValues := range []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1e2, 1e3, 1e4} {
+		for _, numValues := range []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 17, 1e2, 1e3, 1e4} {
 			t.Run(fmt.Sprintf("%s/N=%d", typ, numValues), func(t *testing.T) {
 				testDictionary(t, typ, numValues)
 			})
@@ -69,6 +69,20 @@ func testDictionary(t *testing.T, typ parquet.Type, numValues int) {
 			}
 		}
 
+		// second insert is a no-op since all the values are already in the dictionary
+		lastDictLen := dict.Len()
+		dict.Insert(indexes[i:j], values[i:j])
+
+		if dict.Len() != lastDictLen {
+			for k, index := range indexes[i:j] {
+				if index >= int32(len(mapping)) {
+					t.Log(values[i+k])
+				}
+			}
+
+			t.Fatalf("%d values were inserted on the second pass", dict.Len()-len(mapping))
+		}
+
 		r.Shuffle(j-i, func(a, b int) {
 			indexes[a+i], indexes[b+i] = indexes[b+i], indexes[a+i]
 		})
diff --git a/sparse/gather_test.go b/sparse/gather_test.go
index 7ed1ed7..11dfeb6 100644
--- a/sparse/gather_test.go
+++ b/sparse/gather_test.go
@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math"
+	"strconv"
 	"testing"
 	"time"
 	"unsafe"
@@ -110,6 +111,35 @@ func ExampleGatherUint128() {
 	// points[9].Y = 18
 }
 
+func ExampleGatherString() {
+	buf := make([][2]string, 10)
+	dst := make([]string, 10)
+	src := sparse.UnsafeStringArray(unsafe.Pointer(&buf[0][1]), len(buf), unsafe.Sizeof(buf[0]))
+
+	for i := range buf {
+		buf[i][0] = "-"
+		buf[i][1] = strconv.Itoa(i)
+	}
+
+	n := sparse.GatherString(dst, src)
+
+	for i, v := range dst[:n] {
+		fmt.Printf("points[%d].Y = %v\n", i, v)
+	}
+
+	// Output:
+	// points[0].Y = 0
+	// points[1].Y = 1
+	// points[2].Y = 2
+	// points[3].Y = 3
+	// points[4].Y = 4
+	// points[5].Y = 5
+	// points[6].Y = 6
+	// points[7].Y = 7
+	// points[8].Y = 8
+	// points[9].Y = 9
+}
+
 func TestGatherUint32(t *testing.T) {
 	type point2D struct{ X, Y uint32 }