workload: allow BatchedTuples to contain duplicates

Previously the BatchedTuples used to load initial data for a workload
disallowed duplicate keys violating uniqueness constraints. Some of the
workload key generators may inevitably generate duplicates when hashing keys
because of the relatively small key space of 64-bit integers.

This commit updates the BatchedTuples type so that the workload can indicate
the possible existence of duplicates, and updates the data loader to use INSERT
... ON CONFLICT DO NOTHING queries to perform the insertion, effectively
ignoring the duplicates.

Author: jbowens

pkg/workload/kv/kv.go

@@ -428,6 +428,13 @@ func (w *kv) Tables() []workload.Table {
 		const batchSize = 1000
 		table.InitialRows = workload.BatchedTuples{
 			NumBatches: (w.insertCount + batchSize - 1) / batchSize,
+			// If the key sequence is not sequential, duplicates are possible.
+			// The zipfian distribution produces duplicates by design, and the
+			// hash key mapper can also produce duplicates at larger insert
+			// counts (it's at least inevitable at ~1b rows). Marking that the
+			// keys may contain duplicates will cause the data loader to use
+			// INSERT ... ON CONFLICT DO NOTHING statements.
+			MayContainDuplicates: !w.sequential,
 			FillBatch: func(batchIdx int, cb coldata.Batch, a *bufalloc.ByteAllocator) {
 				rowBegin, rowEnd := batchIdx*batchSize, (batchIdx+1)*batchSize
 				if rowEnd > w.insertCount {

pkg/workload/workload.go

@@ -46,7 +46,7 @@ type Generator interface {
 func SupportsFixtures(gen Generator) bool {
 	tt := gen.Tables()
 	for _, t := range tt {
-		if t.InitialRows.FillBatch == nil {
+		if t.InitialRows.FillBatch == nil || t.InitialRows.MayContainDuplicates {
 			return false
 		}
 	}
@@ -205,6 +205,10 @@ func (t Table) GetResolvedName() tree.TableName {
 type BatchedTuples struct {
 	// NumBatches is the number of batches of tuples.
 	NumBatches int
+	// MayContainDuplicates is a flag indicating whether the tuples may contain
+	// keys that violate uniqueness constraints. If true, the data loader will
+	// use INSERT ... ON CONFLICT DO NOTHING statements.
+	MayContainDuplicates bool
 	// FillBatch is a function to deterministically compute a columnar-batch of
 	// tuples given its index.
 	//

pkg/workload/workloadsql/dataload.go

@@ -131,6 +131,9 @@ func (l InsertsDataLoader) InitialDataLoad(
 				var numRows int
 				flush := func() error {
 					if len(params) > 0 {
+						if table.InitialRows.MayContainDuplicates {
+							fmt.Fprint(&insertStmtBuf, ` ON CONFLICT DO NOTHING`)
+						}
 						insertStmt := insertStmtBuf.String()
 						if _, err := db.ExecContext(gCtx, insertStmt, params...); err != nil {
 							return errors.Wrapf(err, "failed insert into %s", tableName.String())

pkg/workload/ycsb/ycsb.go

@@ -359,6 +359,12 @@ func (g *ycsb) Tables() []workload.Table {
 		const batchSize = 1000
 		usertable.InitialRows = workload.BatchedTuples{
 			NumBatches: (g.insertCount + batchSize - 1) / batchSize,
+			// If the key sequence is hashed, duplicates are possible. Hash
+			// collisions are inevitable at large insert counts (they're at
+			// least inevitable at ~1b rows). Marking that the keys may contain
+			// duplicates will cause the data loader to use INSERT ... ON
+			// CONFLICT DO NOTHING statements.
+			MayContainDuplicates: !g.insertHash,
 			FillBatch: func(batchIdx int, cb coldata.Batch, _ *bufalloc.ByteAllocator) {
 				rowBegin, rowEnd := batchIdx*batchSize, (batchIdx+1)*batchSize
 				if rowEnd > g.insertCount {