workload: use same key sequence during KV --insert-count

When inserting an initial --insert-count keys, use the configured key schema to
insert the keys. This ensures that a subsequent call passing an appropriate
--write-seq flag will read the keys inserted through --insert-count.

Fix #107874.
Epic: none
Release note: none

Author: jbowens

pkg/workload/kv/kv.go

@@ -389,12 +389,6 @@ func splitFinder(i, splits int, r keyRange, k keyTransformer) interface{} {
 	return k.getKey(splitPoint)
 }
 
-func insertCountKey(idx, count int64, kr keyRange) int64 {
-	stride := kr.max/(count+1) - kr.min/(count+1)
-	key := kr.min + (idx+1)*stride
-	return key
-}
-
 // Tables implements the Generator interface.
 func (w *kv) Tables() []workload.Table {
 	// Tables should only run on initialized workload, safe to call create without
@@ -429,6 +423,9 @@ func (w *kv) Tables() []workload.Table {
 		table.InitialRows = workload.BatchedTuples{
 			NumBatches: (w.insertCount + batchSize - 1) / batchSize,
 			FillBatch: func(batchIdx int, cb coldata.Batch, a *bufalloc.ByteAllocator) {
+				// Grab a new state for each batch, under the assumption that
+				// FillBatch may be called concurrently.
+				ks := kg.newState()
 				rowBegin, rowEnd := batchIdx*batchSize, (batchIdx+1)*batchSize
 				if rowEnd > w.insertCount {
 					rowEnd = w.insertCount
@@ -444,11 +441,11 @@ func (w *kv) Tables() []workload.Table {
 				{
 					seq := rowBegin
 					kg.transformer.fillColumnBatch(cb, a, func() (s int64, ok bool) {
-						if seq < rowEnd {
-							seq++
-							return insertCountKey(int64(seq-1), int64(w.insertCount), kg.kr), true
+						if seq >= rowEnd {
+							return 0, false
 						}
-						return 0, false
+						seq++
+						return ks.mapKey.mapKey(int64(seq)), true
 					})
 				}