diff --git a/flow/connectors/mysql/qrep.go b/flow/connectors/mysql/qrep.go index 8c0e5532c..64cf5d16d 100644 --- a/flow/connectors/mysql/qrep.go +++ b/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) diff --git a/flow/connectors/mysql/qrep_partition.go b/flow/connectors/mysql/qrep_partition.go index b285088dd..8a53428db 100644 --- a/flow/connectors/mysql/qrep_partition.go +++ b/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,234 @@ func bigIntToUUID(n *big.Int, casing hexCasing) (string, error) { } return s, nil } + +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) + fetchPrevRealKey( + 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 { + // The interpolated midpoint can overshoot every key in the range when + // keys occupy a narrow slice of the character space. So also probe + // backwards before declaring the partition unsplittable. + k, found, err = prober.fetchPrevRealKey(ctx, tableName, quotedCol, mid, p.start, p.end) + if err != nil { + return nil, fmt.Errorf("failed to fetch prev 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, utils.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)) + return c.fetchRealKey(ctx, query) +} + +// fetchPrevRealKey returns the largest real value of the watermark column that +// is below 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) fetchPrevRealKey( + 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 DESC LIMIT 1", + quotedCol, tableName, mysql.Escape(midpoint), mysql.Escape(start), mysql.Escape(end)) + return c.fetchRealKey(ctx, query) +} + +func (c *MySqlConnector) fetchRealKey(ctx context.Context, query string) (string, bool, error) { + 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 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 +} diff --git a/flow/connectors/mysql/qrep_partition_test.go b/flow/connectors/mysql/qrep_partition_test.go index 501f36c75..64eff14d6 100644 --- a/flow/connectors/mysql/qrep_partition_test.go +++ b/flow/connectors/mysql/qrep_partition_test.go @@ -1,13 +1,22 @@ package connmysql import ( + "container/heap" + "context" + "fmt" + "log/slog" "regexp" + "slices" + "strings" "testing" + "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.temporal.io/sdk/log" "github.com/PeerDB-io/peerdb/flow/generated/protos" + "github.com/PeerDB-io/peerdb/flow/pkg/common" ) func TestUUIDToBigIntRoundTrip(t *testing.T) { @@ -165,3 +174,266 @@ func TestBuildUuidStringPartitionsInvalid(t *testing.T) { }) } } + +func TestStringMidpoint(t *testing.T) { + cases := []struct { + name string + s1 string + s2 string + expected string + }{ + {"identical", "abc", "abc", "abc"}, + {"empty bounds", "", "", ""}, + {"with spaces", " m e", " o ghi", " n fDDOOO"}, + {"numeric string", "111", "999", "555"}, + {"with shared prefix", "prefix-hello", "prefix-world", "prefix-p;@=:"}, + {"max length", "qwertyuiop", "zxcvbnmlkj", "vHdtktB;"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + mid := stringMidpoint(tc.s1, tc.s2) + assert.Equal(t, tc.expected, mid) + assert.LessOrEqual(t, tc.s1, mid) + assert.LessOrEqual(t, mid, tc.s2) + }) + } +} + +func TestBase95RoundTrip(t *testing.T) { + trimRightSpaces := func(s string) string { + return regexp.MustCompile(` +$`).ReplaceAllString(s, "") + } + + for _, s := range []string{"", "a", "ant", "cat", "~~~~~~~~", "555", "Mixed42!", "longer_than_8_characters"} { + decoded := base95IntegerToString(stringToBase95Integer(s)) + expected := s + if len(expected) > 8 { + expected = expected[0:8] + } + assert.Equal(t, expected, trimRightSpaces(decoded)) + } +} + +func TestStringPartitionHeapPopsLargest(t *testing.T) { + h := &stringPartitionHeap{} + heap.Init(h) + for _, rows := range []uint64{5, 100, 1, 42, 7} { + heap.Push(h, stringPartitionEntry{rows: rows}) + } + var popped []uint64 + for h.Len() > 0 { + popped = append(popped, heap.Pop(h).(stringPartitionEntry).rows) + } + assert.Equal(t, []uint64{100, 42, 7, 5, 1}, popped) +} + +type fakeRangeProber struct { + keys []string + less func(a string, b string) bool +} + +func newFakeRangeProber(keys []string, less func(a string, b string) bool) *fakeRangeProber { + sorted := slices.Clone(keys) + slices.SortFunc(sorted, func(a string, b string) int { + switch { + case less(a, b): + return -1 + case less(b, a): + return 1 + default: + return 0 + } + }) + return &fakeRangeProber{keys: sorted, less: less} +} + +func (f *fakeRangeProber) estimateRowsInRange( + _ context.Context, _ string, _ string, start string, end string, +) (uint64, error) { + var n uint64 + for _, k := range f.keys { + if !f.less(k, start) && f.less(k, end) { // start <= k < end + n++ + } + } + return n, nil +} + +func (f *fakeRangeProber) fetchNextRealKey( + _ context.Context, _ string, _ string, midpoint string, start string, end string, +) (string, bool, error) { + for _, k := range f.keys { + if f.less(k, midpoint) { // k < midpoint + continue + } + if !f.less(start, k) { // k <= start + continue + } + if !f.less(k, end) { // k >= end + continue + } + return k, true, nil + } + return "", false, nil +} + +func (f *fakeRangeProber) fetchPrevRealKey( + _ context.Context, _ string, _ string, midpoint string, start string, end string, +) (string, bool, error) { + for i := len(f.keys) - 1; i >= 0; i-- { + k := f.keys[i] + if !f.less(k, midpoint) { // k >= midpoint + continue + } + if !f.less(start, k) { // k <= start + continue + } + if !f.less(k, end) { // k >= end + continue + } + return k, true, nil + } + return "", false, nil +} + +// binaryLess mimics a *_bin collation +func binaryLess(a string, b string) bool { return a < b } + +// caseInsensitiveLess mimics a *_ci collation +func caseInsensitiveLess(a string, b string) bool { + la, lb := strings.ToLower(a), strings.ToLower(b) + if la != lb { + return la < lb + } + return a < b +} + +func runAdaptiveStringPartitions(t *testing.T, keys []string, less func(a string, b string) bool, numPartitions int64) []*protos.QRepPartition { + t.Helper() + planner := newFakeRangeProber(keys, less) + minVal, maxVal := planner.keys[0], planner.keys[len(planner.keys)-1] + table := &common.QualifiedTable{Namespace: "db", Table: "t"} + partitions, err := buildAdaptiveStringPartitions( + t.Context(), planner, log.NewStructuredLogger(slog.Default()), + table, "id", minVal, maxVal, numPartitions) + require.NoError(t, err) + return partitions +} + +func verifyFullCoverage(t *testing.T, partitions []*protos.QRepPartition, keys []string, less func(a string, b string) bool) { + t.Helper() + require.NotEmpty(t, partitions) + + // verify one end-inclusive partition exists + count := 0 + for _, p := range partitions { + if p.GetRange().GetStringRange().EndInclusive { + count++ + } + } + require.Equal(t, 1, count) + + // verify each key belongs to exactly one partition + for _, key := range keys { + matched := -1 + for i, p := range partitions { + sr := p.GetRange().GetStringRange() + require.NotNil(t, sr) + var inRange bool + if sr.EndInclusive { + inRange = !less(key, sr.Start) && !less(sr.End, key) // start <= key <= end + } else { + inRange = !less(key, sr.Start) && less(key, sr.End) // start <= key < end + } + if inRange { + require.Equal(t, -1, matched) + matched = i + } + } + require.NotEqual(t, -1, matched) + } +} + +func TestBuildAdaptiveStringPartitions_WellDistributed(t *testing.T) { + var keys []string + for c := byte('a'); c <= byte('z'); c++ { + keys = append(keys, fmt.Sprintf("%c_value", c)) + } + const maxNumPartitions = 8 + partitions := runAdaptiveStringPartitions(t, keys, binaryLess, maxNumPartitions) + require.Len(t, partitions, maxNumPartitions) + verifyFullCoverage(t, partitions, keys, binaryLess) +} + +func TestBuildAdaptiveStringPartitions_NumKeysLessThanNumPartitions(t *testing.T) { + keys := []string{"alpha", "bravo", "charlie", "delta"} + partitions := runAdaptiveStringPartitions(t, keys, binaryLess, 8) + // [alpha, bravo), [bravo, charlie), [charlie, delta] + require.Len(t, partitions, 3) + verifyFullCoverage(t, partitions, keys, binaryLess) +} + +func TestBuildAdaptiveStringPartitions_LowCardinality(t *testing.T) { + keys := []string{"aaa", "aaa", "aaa", "ooo", "ooo", "zzz", "zzz", "zzz", "zzz"} + distinct := len(slices.Compact(slices.Clone(keys))) + partitions := runAdaptiveStringPartitions(t, keys, binaryLess, 8) + require.LessOrEqual(t, len(partitions), distinct-1) + require.GreaterOrEqual(t, len(partitions), 1) + verifyFullCoverage(t, partitions, keys, binaryLess) +} + +func TestBuildAdaptiveStringPartitions_AllIdentical(t *testing.T) { + keys := []string{"same", "same", "same", "same"} + partitions := runAdaptiveStringPartitions(t, keys, binaryLess, 8) + require.Len(t, partitions, 1) + verifyFullCoverage(t, partitions, keys, binaryLess) +} + +func TestBuildAdaptiveStringPartitions_Unicode(t *testing.T) { + keys := []string{ + "αlpha", "βeta", "γamma", "日本語", "中文", "🍕pizza", "🍟fries", "smiley😀", "café", "naïve", + } + partitions := runAdaptiveStringPartitions(t, keys, binaryLess, 8) + verifyFullCoverage(t, partitions, keys, binaryLess) +} + +func TestBuildAdaptiveStringPartitions_CaseInsensitiveCollation(t *testing.T) { + keys := []string{ + "Apple", "apple", "BANANA", "banana", "Cherry", "cherry", "Date", "DATE", "elderberry", "Fig", + } + partitions := runAdaptiveStringPartitions(t, keys, caseInsensitiveLess, 8) + verifyFullCoverage(t, partitions, keys, caseInsensitiveLess) +} + +func TestBuildAdaptiveStringPartitions_NarrowRange(t *testing.T) { + keys := make([]string, 0, 100) + for i := 1; i <= 100; i++ { + keys = append(keys, fmt.Sprintf("key_%04d", i)) + } + + // demonstrate fetchNextRealKey would not return a valid key + // (e.g. "key_00_" sorts after "key_0099") so splitting these + // keys relies on the fetchPrevRealKey fallback + require.Equal(t, "key_00_`", stringMidpoint("key_0001", "key_0100")) + + partitions := runAdaptiveStringPartitions(t, keys, binaryLess, 8) + require.Len(t, partitions, 8) + verifyFullCoverage(t, partitions, keys, binaryLess) +} + +func TestBuildAdaptiveStringPartitions_MixedCase_UUID(t *testing.T) { + keys := make([]string, 0, 100) + for i := 1; i <= 100; i++ { + id, err := uuid.NewUUID() + require.NoError(t, err) + if i%2 == 0 { + keys = append(keys, id.String()) + } else { + keys = append(keys, strings.ToUpper(id.String())) + } + } + const numPartitions = 5 + partitions := runAdaptiveStringPartitions(t, keys, binaryLess, numPartitions) + require.Len(t, partitions, numPartitions) + verifyFullCoverage(t, partitions, keys, binaryLess) +} diff --git a/flow/e2e/clickhouse_mysql_test.go b/flow/e2e/clickhouse_mysql_test.go index 56f49e214..e6e8425fc 100644 --- a/flow/e2e/clickhouse_mysql_test.go +++ b/flow/e2e/clickhouse_mysql_test.go @@ -5,6 +5,7 @@ import ( "encoding/hex" "fmt" "math" + rand "math/rand/v2" "strconv" "strings" "testing" @@ -2203,7 +2204,7 @@ func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_UUID_Parallel_Snapshot( RequireEnvCanceled(s.t, env) } -func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_Arbitrary_FullTable() { +func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_Arbitrary_Parallel_Snapshot() { if _, ok := s.source.(*MySqlSource); !ok { s.t.Skip("only applies to mysql") } @@ -2213,13 +2214,28 @@ func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_Arbitrary_FullTable() { dstTable := "test_string_pk_non_uuid_dst" const numRows = 100 - const numPartitions = 8 + const numPartitions = 10 require.NoError(s.t, s.source.Exec(s.t.Context(), fmt.Sprintf("CREATE TABLE %s (id VARCHAR(50) PRIMARY KEY, val INT NOT NULL)", srcFullName))) - for i := 1; i <= numRows; i++ { - require.NoError(s.t, s.source.Exec(s.t.Context(), - fmt.Sprintf("INSERT INTO %s (id, val) VALUES ('key_%04d', %d)", srcFullName, i, i))) + seen := make(map[string]struct{}, numRows) + values := make([]string, 0, numRows) + for len(seen) < numRows { + keyBytes := make([]byte, 10) + for i := range keyBytes { + keyBytes[i] = byte('a' + rand.IntN(26)) //nolint:gosec // test data + } + key := string(keyBytes) + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + values = append(values, fmt.Sprintf("('%s', %d)", key, rand.IntN(100))) //nolint:gosec // test data } + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf("INSERT INTO %s (id, val) VALUES %s", srcFullName, strings.Join(values, ",")))) + + // to make stats more accurate on a freshly generated table + require.NoError(s.t, s.source.Exec(s.t.Context(), "ANALYZE TABLE "+srcFullName)) tableMappings := TableMappings(s, srcTable, dstTable) // a String column can't be used directly as a Distributed sharding key (ClickHouse @@ -2259,7 +2275,7 @@ func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_Arbitrary_FullTable() { partitionCount++ } require.NoError(s.t, partitionRows.Err()) - require.EqualValues(s.t, 1, partitionCount) + require.EqualValues(s.t, numPartitions, partitionCount) require.EqualValues(s.t, numRows, totalRows) env.Cancel(s.t.Context())