support arbitrary string parallel snapshotting

Author: jgao54

flow/connectors/mysql/qrep.go

@@ -164,8 +164,14 @@ func (c *MySqlConnector) GetQRepPartitions(
 			}
 			partitionHelper.AddPartitions(uuidPartitions)
 		} else {
-			c.logger.Info("string watermark column is not uuid, falling back to full table partition")
-			return utils.FullTablePartition(), nil
+			stringPartitions, err := buildAdaptiveStringPartitions(
+				ctx, c, c.logger, parsedWatermarkTable, config.WatermarkColumn, start, end, numPartitions)
+			if err != nil {
+				c.logger.Warn("failed to build adaptive string partitions, falling back to full table partition",
+					slog.Any("error", err))
+				return utils.FullTablePartition(), nil
+			}
+			partitionHelper.AddPartitions(stringPartitions)
 		}
 	} else if err := partitionHelper.AddPartitionsWithRange(val1.Value(), val2.Value(), numPartitions); err != nil {
 		return nil, fmt.Errorf("failed to add partitions: %w", err)

flow/connectors/mysql/qrep_partition.go

@@ -1,15 +1,21 @@
 package connmysql
 
 import (
+	"container/heap"
+	"context"
 	"fmt"
+	"log/slog"
 	"math/big"
 	"regexp"
 	"strings"
 
+	"github.com/go-mysql-org/go-mysql/mysql"
 	"github.com/google/uuid"
+	"go.temporal.io/sdk/log"
 
 	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	"github.com/PeerDB-io/peerdb/flow/pkg/common"
 	"github.com/PeerDB-io/peerdb/flow/shared"
 )
 
@@ -100,3 +106,220 @@ func bigIntToUUID(n *big.Int, casing hexCasing) (string, error) {
 	}
 	return s, nil
 }
+
+func createStringPartition(start string, end string, endInclusive bool) *protos.QRepPartition {
+	return &protos.QRepPartition{
+		PartitionId: uuid.NewString(),
+		Range: &protos.PartitionRange{
+			Range: &protos.PartitionRange_StringRange{
+				StringRange: &protos.StringPartitionRange{
+					Start:        start,
+					End:          end,
+					EndInclusive: endInclusive,
+				},
+			},
+		},
+	}
+}
+
+const (
+	base95Min   = ' ' // 0x20, lowest printable ASCII -> digit 0
+	base95Max   = '~' // 0x7E, highest printable ASCII -> digit 94
+	base95Radix = base95Max - base95Min + 1
+	base95Width = 8 // 95^8 to fit in an uint64
+)
+
+func stringMidpoint(s1 string, s2 string) string {
+	i := 0
+	for i < len(s1) && i < len(s2) && s1[i] == s2[i] {
+		i++
+	}
+	sharedPrefix := s1[:i]
+	s1, s2 = s1[i:], s2[i:]
+	mid := (stringToBase95Integer(s1) + stringToBase95Integer(s2)) / 2
+	return strings.TrimRight(sharedPrefix+base95IntegerToString(mid), " ")
+}
+
+func stringToBase95Integer(s string) uint64 {
+	if s == "" {
+		return 0
+	}
+	var res uint64
+	for i := range base95Width {
+		var digit uint64
+		if i < len(s) {
+			ch := s[i]
+			switch {
+			case ch < base95Min:
+				ch = base95Min
+			case ch > base95Max:
+				ch = base95Max
+			}
+			digit = uint64(ch - base95Min)
+		}
+		res = res*base95Radix + digit
+	}
+	return res
+}
+
+func base95IntegerToString(n uint64) string {
+	digits := make([]byte, base95Width)
+	for k := base95Width - 1; k >= 0; k-- {
+		digits[k] = base95Min + byte(n%base95Radix)
+		n /= base95Radix
+	}
+	return string(digits)
+}
+
+type stringPartitionEntry struct {
+	start string
+	end   string
+	rows  uint64
+}
+
+type stringPartitionHeap []stringPartitionEntry
+
+func (h *stringPartitionHeap) Len() int           { return len(*h) }
+func (h *stringPartitionHeap) Less(i, j int) bool { return (*h)[i].rows > (*h)[j].rows }
+func (h *stringPartitionHeap) Swap(i, j int)      { (*h)[i], (*h)[j] = (*h)[j], (*h)[i] }
+
+func (h *stringPartitionHeap) Push(x any) { *h = append(*h, x.(stringPartitionEntry)) }
+
+func (h *stringPartitionHeap) Pop() any {
+	old := *h
+	n := len(old)
+	item := old[n-1]
+	*h = old[:n-1]
+	return item
+}
+
+// interface for unit-testing
+type rangeProber interface {
+	estimateRowsInRange(ctx context.Context, tableName string, quotedCol string, start string, end string) (uint64, error)
+	fetchNextRealKey(
+		ctx context.Context, tableName string, quotedCol string, midpoint string, start string, end string,
+	) (string, bool, error)
+}
+
+// buildAdaptiveStringPartitions splits an arbitrary string watermark column
+// into at most numPartitions partitions using midpoint bisection guided
+// by the query planner's row estimates. It starts from a single [minVal, maxVal]
+// partition and repeatedly splits the largest partition, until it reaches
+// numPartitions or runs out of splittable partitions.
+func buildAdaptiveStringPartitions(
+	ctx context.Context,
+	prober rangeProber,
+	logger log.Logger,
+	table *common.QualifiedTable,
+	watermarkColumn string,
+	minVal string,
+	maxVal string,
+	numPartitions int64,
+) ([]*protos.QRepPartition, error) {
+	tableName := table.MySQL()
+	quotedCol := common.QuoteMySQLIdentifier(watermarkColumn)
+
+	totalRows, err := prober.estimateRowsInRange(ctx, tableName, quotedCol, minVal, maxVal)
+	if err != nil {
+		return nil, fmt.Errorf("failed to estimate rows for [%s, %s]: %w", minVal, maxVal, err)
+	}
+	h := &stringPartitionHeap{{start: minVal, end: maxVal, rows: totalRows}}
+	heap.Init(h)
+
+	var outputs []stringPartitionEntry
+	for int64(len(outputs)+h.Len()) < numPartitions && h.Len() > 0 {
+		p := heap.Pop(h).(stringPartitionEntry)
+
+		if p.start == p.end {
+			outputs = append(outputs, p)
+			continue
+		}
+
+		mid := stringMidpoint(p.start, p.end)
+		k, found, err := prober.fetchNextRealKey(ctx, tableName, quotedCol, mid, p.start, p.end)
+		if err != nil {
+			return nil, fmt.Errorf("failed to fetch next real key for [%s, %s]: %w", p.start, p.end, err)
+		}
+		if !found {
+			outputs = append(outputs, p)
+			continue
+		}
+
+		leftRows, err := prober.estimateRowsInRange(ctx, tableName, quotedCol, p.start, k)
+		if err != nil {
+			return nil, fmt.Errorf("failed to estimate rows for [%s, %s]: %w", p.start, k, err)
+		}
+		rightRows, err := prober.estimateRowsInRange(ctx, tableName, quotedCol, k, p.end)
+		if err != nil {
+			return nil, fmt.Errorf("failed to estimate rows for [%s, %s]: %w", k, p.end, err)
+		}
+		heap.Push(h, stringPartitionEntry{start: p.start, end: k, rows: leftRows})
+		heap.Push(h, stringPartitionEntry{start: k, end: p.end, rows: rightRows})
+	}
+
+	for h.Len() > 0 {
+		outputs = append(outputs, heap.Pop(h).(stringPartitionEntry))
+	}
+
+	partitions := make([]*protos.QRepPartition, 0, len(outputs))
+	for _, p := range outputs {
+		partitions = append(partitions, createStringPartition(p.start, p.end, p.end == maxVal))
+	}
+	logger.Info("[mysql] built adaptive string partitions",
+		slog.Int64("targetNumPartitions", numPartitions),
+		slog.Int("numPartitions", len(partitions)))
+
+	return partitions, nil
+}
+
+// fetchNextRealKey returns the smallest real value of the watermark column that
+// is at or past the interpolated midpoint and strictly inside (start, end). The
+// bounds are enforced by MySQL using its column's collation. Return false when
+// no such key exists.
+func (c *MySqlConnector) fetchNextRealKey(
+	ctx context.Context, tableName string, quotedCol string, midpoint string, start string, end string,
+) (string, bool, error) {
+	query := fmt.Sprintf(
+		"SELECT %[1]s FROM %[2]s WHERE %[1]s >= '%[3]s' AND %[1]s > '%[4]s' AND %[1]s < '%[5]s' ORDER BY %[1]s LIMIT 1",
+		quotedCol, tableName, mysql.Escape(midpoint), mysql.Escape(start), mysql.Escape(end))
+	rs, err := c.Execute(ctx, query)
+	if err != nil {
+		return "", false, err
+	}
+	defer rs.Close()
+	if rs.RowNumber() == 0 {
+		return "", false, nil
+	}
+	key, err := rs.GetString(0, 0)
+	if err != nil {
+		return "", false, fmt.Errorf("failed to read next real key: %w", err)
+	}
+	// GetString is zero-copy over the resultset's row buffer, which rs.Close()
+	// recycles into a pool where the next query's response overwrites it; the
+	// key is stored in the heap and outlives the function, so it must own its bytes
+	return strings.Clone(key), true, nil
+}
+
+// estimateRowsInRange asks the query planner for the estimated number of rows in
+// [start, end). The estimate is best-effort: a NULL or zero estimate is returned
+// as 0 rather than an error, so the caller still treats the range as a real partition.
+func (c *MySqlConnector) estimateRowsInRange(
+	ctx context.Context, tableName string, quotedCol string, start string, end string,
+) (uint64, error) {
+	query := fmt.Sprintf(
+		"EXPLAIN FORMAT=TRADITIONAL SELECT 1 FROM %[1]s WHERE %[2]s >= '%[3]s' AND %[2]s < '%[4]s'",
+		tableName, quotedCol, mysql.Escape(start), mysql.Escape(end))
+	rs, err := c.Execute(ctx, query)
+	if err != nil {
+		return 0, err
+	}
+	defer rs.Close()
+	if rs.RowNumber() == 0 {
+		return 0, fmt.Errorf("no resultset for table %s", tableName)
+	}
+	rows, err := rs.GetUintByName(0, "rows")
+	if err != nil {
+		return 0, err
+	}
+	return rows, nil
+}