diff --git a/sstable/colblk/data_block.go b/sstable/colblk/data_block.go index 33e1fcf383..1d252dbd2e 100644 --- a/sstable/colblk/data_block.go +++ b/sstable/colblk/data_block.go @@ -662,8 +662,8 @@ func (i *DataBlockIter) Next() *base.InternalKV { Trailer: base.InternalKeyTrailer(i.r.trailers.At(i.row)), } // Inline i.r.values.At(row). - startOffset := i.r.values.offsets.At(i.row) - v := unsafe.Slice((*byte)(i.r.values.ptr(startOffset)), i.r.values.offsets.At(i.row+1)-startOffset) + startOffset, endOffset := i.r.values.offsets.At2(i.row) + v := unsafe.Slice((*byte)(i.r.values.ptr(startOffset)), endOffset-startOffset) if i.r.isValueExternal.At(i.row) { i.kv.V = i.getLazyValue(v) } else { diff --git a/sstable/colblk/prefix_bytes.go b/sstable/colblk/prefix_bytes.go index 225b6657e5..638595cff6 100644 --- a/sstable/colblk/prefix_bytes.go +++ b/sstable/colblk/prefix_bytes.go @@ -255,8 +255,8 @@ type PrefixBytesIter struct { func (b *PrefixBytes) SetAt(it *PrefixBytesIter, i int) { // Determine the offset and length of the bundle prefix. bundleOffsetIndex := b.bundleOffsetIndexForRow(i) - bundleOffsetStart := b.rawBytes.offsets.At(bundleOffsetIndex) - it.bundlePrefixLen = b.rawBytes.offsets.At(bundleOffsetIndex+1) - bundleOffsetStart + bundleOffsetStart, bundleOffsetEnd := b.rawBytes.offsets.At2(bundleOffsetIndex) + it.bundlePrefixLen = bundleOffsetEnd - bundleOffsetStart // Determine the offset and length of the row's individual suffix. it.offsetIndex = b.rowSuffixIndex(i) @@ -265,9 +265,9 @@ func (b *PrefixBytes) SetAt(it *PrefixBytesIter, i int) { // this recomputation? The expected case is non-duplicate keys, so it may // not be worthwhile. rowSuffixStart, rowSuffixEnd := b.rowSuffixOffsets(i, it.offsetIndex) + rowSuffixLen := rowSuffixEnd - rowSuffixStart - // Grow the size of the iterator's buffer if necessary. - it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen)+int(rowSuffixEnd-rowSuffixStart)] + it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen+rowSuffixLen)] ptr := unsafe.Pointer(unsafe.SliceData(it.buf)) // Copy the shared key prefix. @@ -281,7 +281,7 @@ func (b *PrefixBytes) SetAt(it *PrefixBytesIter, i int) { memmove( unsafe.Pointer(uintptr(ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)), unsafe.Pointer(uintptr(b.rawBytes.data)+uintptr(rowSuffixStart)), - uintptr(rowSuffixEnd-rowSuffixStart)) + uintptr(rowSuffixLen)) // Set nextBundleOffsetIndex so that a call to SetNext can cheaply determine // whether the next row is in the same bundle. it.nextBundleOffsetIndex = bundleOffsetIndex + (1 << b.bundleShift) + 1 @@ -297,22 +297,21 @@ func (b *PrefixBytes) SetNext(it *PrefixBytesIter) { // If the next row is in the same bundle, we can take a fast path of only // updating the per-row suffix. if it.offsetIndex < it.nextBundleOffsetIndex { - rowSuffixStart := b.rawBytes.offsets.At(it.offsetIndex) - rowSuffixEnd := b.rawBytes.offsets.At(it.offsetIndex + 1) - if rowSuffixStart == rowSuffixEnd { + rowSuffixStart, rowSuffixEnd := b.rawBytes.offsets.At2(it.offsetIndex) + rowSuffixLen := rowSuffixEnd - rowSuffixStart + if rowSuffixLen == 0 { // The start and end offsets are equal, indicating that the key is a // duplicate. Since it's identical to the previous key, there's // nothing left to do, we can leave buf as-is. return } - // Grow the buffer if necessary. - it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen)+int(rowSuffixEnd-rowSuffixStart)] + it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen+rowSuffixLen)] // Copy in the per-row suffix. ptr := unsafe.Pointer(unsafe.SliceData(it.buf)) memmove( unsafe.Pointer(uintptr(ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)), unsafe.Pointer(uintptr(b.rawBytes.data)+uintptr(rowSuffixStart)), - uintptr(rowSuffixEnd-rowSuffixStart)) + uintptr(rowSuffixLen)) return } @@ -320,8 +319,8 @@ func (b *PrefixBytes) SetNext(it *PrefixBytesIter) { // The offsetIndex is currently pointing to the start of the new bundle // prefix. Increment it to point at the start of the new row suffix. it.offsetIndex++ - rowSuffixStart := b.rawBytes.offsets.At(it.offsetIndex) - rowSuffixEnd := b.rawBytes.offsets.At(it.offsetIndex + 1) + rowSuffixStart, rowSuffixEnd := b.rawBytes.offsets.At2(it.offsetIndex) + rowSuffixLen := rowSuffixEnd - rowSuffixStart // Read the offsets of the new bundle prefix and update the index of the // next bundle. @@ -329,8 +328,7 @@ func (b *PrefixBytes) SetNext(it *PrefixBytesIter) { it.bundlePrefixLen = rowSuffixStart - bundlePrefixStart it.nextBundleOffsetIndex = it.offsetIndex + (1 << b.bundleShift) - // Grow the buffer if necessary. - it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen)+int(rowSuffixEnd-rowSuffixStart)] + it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen+rowSuffixLen)] // Copy in the new bundle suffix. ptr := unsafe.Pointer(unsafe.SliceData(it.buf)) memmove( @@ -341,7 +339,7 @@ func (b *PrefixBytes) SetNext(it *PrefixBytesIter) { memmove( unsafe.Pointer(uintptr(ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)), unsafe.Pointer(uintptr(b.rawBytes.data)+uintptr(rowSuffixStart)), - uintptr(rowSuffixEnd-rowSuffixStart)) + uintptr(rowSuffixLen)) } // SharedPrefix return a []byte of the shared prefix that was extracted from @@ -383,8 +381,7 @@ func (b *PrefixBytes) RowSuffix(row int) []byte { func (b *PrefixBytes) rowSuffixOffsets(row, i int) (low uint32, high uint32) { // Retrieve the low and high offsets indicating the start and end of the // row's suffix slice. - low = b.rawBytes.offsets.At(i) - high = b.rawBytes.offsets.At(i + 1) + low, high = b.rawBytes.offsets.At2(i) // If there's a non-empty slice for the row, this row is different than its // predecessor. if low != high { diff --git a/sstable/colblk/unsafe_slice.go b/sstable/colblk/unsafe_slice.go index bc3b3966dd..c8c63e8a69 100644 --- a/sstable/colblk/unsafe_slice.go +++ b/sstable/colblk/unsafe_slice.go @@ -92,23 +92,27 @@ func makeUnsafeUints(base uint64, ptr unsafe.Pointer, width int) UnsafeUints { // At returns the `i`-th element. func (s UnsafeUints) At(i int) uint64 { - // TODO(jackson): Experiment with other alternatives that might be faster - // and avoid switching on the width. - switch s.width { - case 0: - return s.base - case 1: - return s.base + uint64(*(*uint8)(unsafe.Pointer(uintptr(s.ptr) + uintptr(i)))) - case 2: - return s.base + uint64(*(*uint16)(unsafe.Pointer(uintptr(s.ptr) + uintptr(i)<> 16 + } + if s.width <= 1 { + if s.width == 0 { + return 0, 0 + } + v := *(*uint16)(unsafe.Pointer(uintptr(s.ptr) + uintptr(i))) + return uint32(v & 0xFF), uint32(v >> 8) + } + v := *(*uint64)(unsafe.Pointer(uintptr(s.ptr) + uintptr(i)<> 32) +} + // UnsafeBuf provides a buffer without bounds checking. Every buf has a len and // capacity. type UnsafeBuf struct { diff --git a/sstable/colblk/unsafe_slice_test.go b/sstable/colblk/unsafe_slice_test.go index c1ab661bcd..53aa5e5eaf 100644 --- a/sstable/colblk/unsafe_slice_test.go +++ b/sstable/colblk/unsafe_slice_test.go @@ -15,11 +15,56 @@ import ( "golang.org/x/exp/rand" ) +func TestUnsafeUints(t *testing.T) { + rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano()))) + for _, r := range interestingIntRanges { + t.Run(r.ExpectedEncoding.String(), func(t *testing.T) { + for _, rows := range []int{1, 10, 100, 10000} { + t.Run(fmt.Sprint(rows), func(t *testing.T) { + vals := make([]uint64, rows) + for i := range vals { + vals[i] = r.Rand(rng) + } + var ub UintBuilder + ub.Init() + for i := 0; i < rows; i++ { + ub.Set(i, vals[i]) + } + sz := ub.Size(rows, 0) + buf := aligned.ByteSlice(int(sz) + 1 /* trailing padding byte */) + _ = ub.Finish(0, rows, 0, buf) + + uints, _ := DecodeUnsafeUints(buf, 0, rows) + for i := range rows { + if uints.At(i) != vals[i] { + t.Fatalf("mismatch at row %d: got %d, expected %d", i, uints.At(i), vals[i]) + } + } + if encoding := UintEncoding(buf[0]); encoding.Width() <= 4 && !encoding.IsDelta() { + offsets, _ := DecodeUnsafeOffsets(buf, 0, rows) + for i := range rows { + if uint64(offsets.At(i)) != vals[i] { + t.Fatalf("mismatch at row %d: got %d, expected %d", i, uints.At(i), vals[i]) + } + } + for i := 0; i < rows-1; i++ { + a, b := offsets.At2(i) + if uint64(a) != vals[i] || uint64(b) != vals[i+1] { + t.Fatalf("mismatch at row %d: got %d,%d, expected %d,%d", i, a, b, vals[i], vals[i+1]) + } + } + } + }) + } + }) + } +} + func BenchmarkUnsafeUints(b *testing.B) { rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano()))) intRanges := []intRange{ - // const - {Min: 1, Max: 1, ExpectedEncoding: makeUintEncoding(0, true)}, + // zero + {Min: 0, Max: 0, ExpectedEncoding: makeUintEncoding(0, false)}, // 1b {Min: 10, Max: 200, ExpectedEncoding: makeUintEncoding(1, false)}, // 1b,delta @@ -31,10 +76,11 @@ func BenchmarkUnsafeUints(b *testing.B) { // 4b {Min: 0, Max: math.MaxUint32, ExpectedEncoding: makeUintEncoding(4, false)}, // 4b,delta - {Min: 100_000, Max: math.MaxUint32 + 10, ExpectedEncoding: makeUintEncoding(4, true)}, + {Min: 1_000_000_100, Max: math.MaxUint32 + 1_000_000_000, ExpectedEncoding: makeUintEncoding(4, true)}, // 8b {Min: 0, Max: math.MaxUint64, ExpectedEncoding: makeUintEncoding(8, false)}, } + for _, r := range intRanges { benchmarkUnsafeUints(b, rng, 1000, r) }