tpcc: Add list and range partitioned table options

cmd/go-tpc/tpcc.go

@@ -76,6 +76,7 @@ func registerTpcc(root *cobra.Command) {
 	}
 
 	cmd.PersistentFlags().IntVar(&tpccConfig.Parts, "parts", 1, "Number to partition warehouses")
+	cmd.PersistentFlags().IntVar(&tpccConfig.PartitionType, "partition-type", 1, "Partition type (1 - HASH, 2 - RANGE, 3 - LIST (like HASH), 4 - LIST (like RANGE)")
 	cmd.PersistentFlags().IntVar(&tpccConfig.Warehouses, "warehouses", 10, "Number of warehouses")
 	cmd.PersistentFlags().BoolVar(&tpccConfig.CheckAll, "check-all", false, "Run all consistency checks")
 	var cmdPrepare = &cobra.Command{

tpcc/csv.go

@@ -42,7 +42,7 @@ func NewCSVWorkloader(db *sql.DB, cfg *Config) (*CSVWorkLoader, error) {
 		cfg:          cfg,
 		initLoadTime: time.Now().Format(timeFormat),
 		tables:       make(map[string]bool),
-		ddlManager:   newDDLManager(cfg.Parts, cfg.UseFK),
+		ddlManager:   newDDLManager(cfg.Parts, cfg.UseFK, cfg.Warehouses, cfg.PartitionType),
 	}
 
 	var val bool

tpcc/ddl.go

@@ -18,12 +18,14 @@ const (
 )
 
 type ddlManager struct {
-	parts int
-	useFK bool
+	parts         int
+	warehouses    int
+	partitionType int
+	useFK         bool
 }
 
-func newDDLManager(parts int, useFK bool) *ddlManager {
-	return &ddlManager{parts: parts, useFK: useFK}
+func newDDLManager(parts int, useFK bool, warehouses, partitionType int) *ddlManager {
+	return &ddlManager{parts: parts, useFK: useFK, warehouses: warehouses, partitionType: partitionType}
 }
 
 func (w *ddlManager) createTableDDL(ctx context.Context, query string, tableName string) error {
@@ -39,8 +41,65 @@ func (w *ddlManager) appendPartition(query string, partKeys string) string {
 	if w.parts <= 1 {
 		return query
 	}
+	if w.partitionType == PartitionTypeListAsHash {
+		// Generate LIST partitions equivalent with HASH partitions
+		s := fmt.Sprintf("%s\nPARTITION BY LIST (%s)\n(", query, partKeys)
+		for i := 0; i < w.parts; i++ {
+			if i > 0 {
+				s = s + ",\n "
+			}
+			var part string
+			for j := i; j < w.warehouses; j = j + w.parts {
+				if j > i {
+					part = part + ","
+				}
+				part = part + fmt.Sprintf("%d", j+1)
+			}
+			s = fmt.Sprintf("%sPARTITION p%d VALUES IN (%s)", s, i, part)
+		}
+		return s + ")"
+	} else if w.partitionType == PartitionTypeListAsRange {
+		// Generate LIST partitions equivalent with RANGE partitions
+		s := fmt.Sprintf("%s\nPARTITION BY LIST (%s)\n(", query, partKeys)
+		addedWarehouses := 0
+		for i := 0; i < w.parts; i++ {
+			if i > 0 {
+				s = s + ",\n "
+			}
+			warehousesToAdd := w.warehouses - addedWarehouses
+			partsLeft := w.parts - i
+			warehousesPerPartition := warehousesToAdd / partsLeft
+			if (warehousesToAdd % partsLeft) != 0 {
+				warehousesPerPartition++
+			}
+			var part string
+			for j := 0; j < warehousesPerPartition; j++ {
+				if j > 0 {
+					part = part + ","
+				}
+				addedWarehouses++
+				part = part + fmt.Sprintf("%d", addedWarehouses)
+			}
+			s = fmt.Sprintf("%sPARTITION p%d VALUES IN (%s)", s, i, part)
+		}
+		return s + ")"
+	} else if w.partitionType == PartitionTypeRange {
+		// Generate RANGE partitions
+		s := fmt.Sprintf("%s\nPARTITION BY RANGE (%s)\n(", query, partKeys)
+		for i := 0; i < w.parts; i++ {
+			if i > 0 {
+				s = s + ",\n "
+			}
+			warehousesPerPartition := w.warehouses / w.parts
+			if (w.warehouses % w.parts) != 0 {
+				warehousesPerPartition++
+			}
+			s = fmt.Sprintf("%sPARTITION p%d VALUES LESS THAN (%d)", s, i, 1+(i+1)*warehousesPerPartition)
+		}
+		return s + ")"
+	}
 
-	return fmt.Sprintf("%s\n PARTITION BY HASH(%s)\n PARTITIONS %d", query, partKeys, w.parts)
+	return fmt.Sprintf("%s\nPARTITION BY HASH(%s)\nPARTITIONS %d", query, partKeys, w.parts)
 }
 
 // createTables creates tables schema.

tpcc/ddl_test.go

@@ -0,0 +1,108 @@
+package tpcc
+
+import "testing"
+
+func TestAppendPartition(t *testing.T) {
+	ddl := newDDLManager(4, false, 4, PartitionTypeHash)
+	s := ddl.appendPartition("<table definition>", "Id")
+	expected := `<table definition>
+PARTITION BY HASH(Id)
+PARTITIONS 4`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+
+	ddl = newDDLManager(4, false, 4, PartitionTypeRange)
+	s = ddl.appendPartition("<table definition>", "Id")
+	expected = `<table definition>
+PARTITION BY RANGE (Id)
+(PARTITION p0 VALUES LESS THAN (2),
+ PARTITION p1 VALUES LESS THAN (3),
+ PARTITION p2 VALUES LESS THAN (4),
+ PARTITION p3 VALUES LESS THAN (5))`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+
+	ddl = newDDLManager(4, false, 23, PartitionTypeRange)
+	s = ddl.appendPartition("<table definition>", "Id")
+	expected = `<table definition>
+PARTITION BY RANGE (Id)
+(PARTITION p0 VALUES LESS THAN (7),
+ PARTITION p1 VALUES LESS THAN (13),
+ PARTITION p2 VALUES LESS THAN (19),
+ PARTITION p3 VALUES LESS THAN (25))`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+
+	ddl = newDDLManager(4, false, 12, PartitionTypeListAsHash)
+	s = ddl.appendPartition("<table definition>", "Id")
+	expected = `<table definition>
+PARTITION BY LIST (Id)
+(PARTITION p0 VALUES IN (1,5,9),
+ PARTITION p1 VALUES IN (2,6,10),
+ PARTITION p2 VALUES IN (3,7,11),
+ PARTITION p3 VALUES IN (4,8,12))`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+
+	ddl = newDDLManager(3, false, 4, PartitionTypeListAsHash)
+	s = ddl.appendPartition("<table definition>", "Id")
+	expected = `<table definition>
+PARTITION BY LIST (Id)
+(PARTITION p0 VALUES IN (1,4),
+ PARTITION p1 VALUES IN (2),
+ PARTITION p2 VALUES IN (3))`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+
+	ddl = newDDLManager(4, false, 23, PartitionTypeListAsHash)
+	s = ddl.appendPartition("<table definition>", "Id")
+	expected = `<table definition>
+PARTITION BY LIST (Id)
+(PARTITION p0 VALUES IN (1,5,9,13,17,21),
+ PARTITION p1 VALUES IN (2,6,10,14,18,22),
+ PARTITION p2 VALUES IN (3,7,11,15,19,23),
+ PARTITION p3 VALUES IN (4,8,12,16,20))`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+
+	ddl = newDDLManager(4, false, 12, PartitionTypeListAsRange)
+	s = ddl.appendPartition("<table definition>", "Id")
+	expected = `<table definition>
+PARTITION BY LIST (Id)
+(PARTITION p0 VALUES IN (1,2,3),
+ PARTITION p1 VALUES IN (4,5,6),
+ PARTITION p2 VALUES IN (7,8,9),
+ PARTITION p3 VALUES IN (10,11,12))`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+
+	ddl = newDDLManager(3, false, 4, PartitionTypeListAsRange)
+	s = ddl.appendPartition("<table definition>", "Id")
+	expected = `<table definition>
+PARTITION BY LIST (Id)
+(PARTITION p0 VALUES IN (1,2),
+ PARTITION p1 VALUES IN (3),
+ PARTITION p2 VALUES IN (4))`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+
+	ddl = newDDLManager(4, false, 23, PartitionTypeListAsRange)
+	s = ddl.appendPartition("<table definition>", "Id")
+	expected = `<table definition>
+PARTITION BY LIST (Id)
+(PARTITION p0 VALUES IN (1,2,3,4,5,6),
+ PARTITION p1 VALUES IN (7,8,9,10,11,12),
+ PARTITION p2 VALUES IN (13,14,15,16,17,18),
+ PARTITION p3 VALUES IN (19,20,21,22,23))`
+	if s != expected {
+		t.Errorf("got '%s' expected '%s'", s, expected)
+	}
+}

tpcc/workload.go

@@ -44,16 +44,24 @@ type tpccState struct {
 	paymentStmts     map[string]*sql.Stmt
 }
 
+const (
+	PartitionTypeHash = iota + 1
+	PartitionTypeRange
+	PartitionTypeListAsHash
+	PartitionTypeListAsRange
+)
+
 // Config is the configuration for tpcc workload
 type Config struct {
-	DBName     string
-	Threads    int
-	Parts      int
-	Warehouses int
-	UseFK      bool
-	Isolation  int
-	CheckAll   bool
-	NoCheck    bool
+	DBName        string
+	Threads       int
+	Parts         int
+	PartitionType int
+	Warehouses    int
+	UseFK         bool
+	Isolation     int
+	CheckAll      bool
+	NoCheck       bool
 
 	// whether to involve wait times(keying time&thinking time)
 	Wait bool
@@ -98,6 +106,10 @@ func NewWorkloader(db *sql.DB, cfg *Config) (workload.Workloader, error) {
 		panic(fmt.Errorf("number warehouses %d must >= partition %d", cfg.Warehouses, cfg.Parts))
 	}
 
+	if cfg.PartitionType < PartitionTypeHash || cfg.PartitionType > PartitionTypeListAsRange {
+		panic(fmt.Errorf("Unknown partition type %d", cfg.PartitionType))
+	}
+
 	resetMaxLat := func(m *measurement.Measurement) {
 		m.MaxLatency = cfg.MaxMeasureLatency
 	}
@@ -106,7 +118,7 @@ func NewWorkloader(db *sql.DB, cfg *Config) (workload.Workloader, error) {
 		db:                  db,
 		cfg:                 cfg,
 		initLoadTime:        time.Now().Format(timeFormat),
-		ddlManager:          newDDLManager(cfg.Parts, cfg.UseFK),
+		ddlManager:          newDDLManager(cfg.Parts, cfg.UseFK, cfg.Warehouses, cfg.PartitionType),
 		rtMeasurement:       measurement.NewMeasurement(resetMaxLat),
 		waitTimeMeasurement: measurement.NewMeasurement(resetMaxLat),
 	}