Skip to content

MySQL: Improve parallel snapshot #3002

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Amogh-Bharadwaj merged 20 commits into from
Jun 6, 2025
Merged

MySQL: Improve parallel snapshot #3002

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
6931ca4
improve mysql partitioning
Amogh-Bharadwaj Jun 3, 2025
b3842b0
account for timestamp and datetime
Amogh-Bharadwaj Jun 3, 2025
268fcc6
refactor partition.go
Amogh-Bharadwaj Jun 3, 2025
b22dbe0
account for nil
Amogh-Bharadwaj Jun 3, 2025
19e377a
use cross join
Amogh-Bharadwaj Jun 3, 2025
6dffa91
limit 0 and single line
Amogh-Bharadwaj Jun 3, 2025
3b6a1c1
drop qk1 and qk2
Amogh-Bharadwaj Jun 3, 2025
2d2b28b
inline err
Amogh-Bharadwaj Jun 3, 2025
9109920
suggestions
Amogh-Bharadwaj Jun 3, 2025
7a430be
lint
Amogh-Bharadwaj Jun 3, 2025
76ca6b0
e2e and use ? instead
Amogh-Bharadwaj Jun 5, 2025
dc16f9c
fix minval passing and include qrep test for mysql
Amogh-Bharadwaj Jun 5, 2025
f342742
improve qrep test
Amogh-Bharadwaj Jun 5, 2025
89ca42c
fix second insert
Amogh-Bharadwaj Jun 5, 2025
0bb0230
patch tests pt 1
Amogh-Bharadwaj Jun 5, 2025
922d3a2
patch test pt2
Amogh-Bharadwaj Jun 5, 2025
512b9df
change expected len
Amogh-Bharadwaj Jun 5, 2025
2bf3dd0
add another wait for tables
Amogh-Bharadwaj Jun 6, 2025
d2d0b4a
rebase fix
Amogh-Bharadwaj Jun 6, 2025
fd3d611
wire source peer
Amogh-Bharadwaj Jun 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
173 changes: 124 additions & 49 deletions flow/connectors/mysql/qrep.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,28 @@ import (
"github.com/PeerDB-io/peerdb/flow/shared/types"
)

func (c *MySqlConnector) GetDataTypeOfWatermarkColumn(
ctx context.Context,
watermarkTableName string,
watermarkColumn string,
) (types.QValueKind, error) {
if watermarkColumn == "" {
return "", errors.New("watermark column is not specified in the config")
}

query := fmt.Sprintf("SELECT `%s` FROM %s LIMIT 0", watermarkColumn, watermarkTableName)
rs, err := c.Execute(ctx, query)
if err != nil {
return "", fmt.Errorf("failed to execute query for watermark column type: %w", err)
}

if len(rs.Fields) == 0 {
return "", errors.New("no fields returned from select query: " + query)
}

return qkindFromMysql(rs.Fields[0])
}

func (c *MySqlConnector) GetQRepPartitions(
ctx context.Context,
config *protos.QRepConfig,
Expand All @@ -46,7 +68,7 @@ func (c *MySqlConnector) GetQRepPartitions(

whereClause := ""
if last != nil && last.Range != nil {
whereClause = fmt.Sprintf("WHERE %s > $1", quotedWatermarkColumn)
whereClause = fmt.Sprintf("WHERE %s > ?", quotedWatermarkColumn)
}
parsedWatermarkTable, err := utils.ParseSchemaTable(config.WatermarkTable)
if err != nil {
Expand All @@ -64,7 +86,7 @@ func (c *MySqlConnector) GetQRepPartitions(
case *protos.PartitionRange_UintRange:
minVal = lastRange.UintRange.End
case *protos.PartitionRange_TimestampRange:
minVal = lastRange.TimestampRange.End.AsTime()
minVal = lastRange.TimestampRange.End.AsTime().String()
}
c.logger.Info(fmt.Sprintf("count query: %s - minVal: %v", countQuery, minVal))

Expand Down Expand Up @@ -99,62 +121,115 @@ func (c *MySqlConnector) GetQRepPartitions(
if totalRows%numRowsPerPartition != 0 {
numPartitions++
}

watermarkQKind, err := c.GetDataTypeOfWatermarkColumn(ctx, parsedWatermarkTable.MySQL(), config.WatermarkColumn)
if err != nil {
return nil, fmt.Errorf("failed to get data type of watermark column %s: %w", config.WatermarkColumn, err)
}

c.logger.Info(fmt.Sprintf("total rows: %d, num partitions: %d, num rows per partition: %d",
totalRows, numPartitions, numRowsPerPartition))
var rs *mysql.Result
if minVal != nil {
// Query to get partitions using window functions
partitionsQuery := fmt.Sprintf(
`SELECT bucket, MIN(%[2]s) AS start, MAX(%[2]s) AS end
FROM (
SELECT NTILE(%[1]d) OVER (ORDER BY %[2]s) AS bucket, %[2]s
FROM %[3]s WHERE %[2]s > $1
) AS subquery

switch watermarkQKind {
case types.QValueKindInt8, types.QValueKindInt16, types.QValueKindInt32, types.QValueKindInt64,
types.QValueKindUInt8, types.QValueKindUInt16, types.QValueKindUInt32, types.QValueKindUInt64:
if minVal != nil {
partitionsQuery := fmt.Sprintf(
`WITH stats AS (
SELECT MIN(%[2]s) AS min_watermark,
1.0 * (MAX(%[2]s) - MIN(%[2]s)) / (%[1]d) AS range_size
FROM %[3]s WHERE %[2]s > ?
)
SELECT FLOOR((w.%[2]s - s.min_watermark) / s.range_size) AS bucket,
MIN(w.%[2]s) AS start, MAX(w.%[2]s) AS end
FROM %[3]s AS w
CROSS JOIN stats AS s
WHERE w.%[2]s > ?
GROUP BY bucket
ORDER BY start`,
numPartitions,
quotedWatermarkColumn,
parsedWatermarkTable.MySQL(),
)
c.logger.Info("partitions query", slog.String("query", partitionsQuery), slog.Any("minVal", minVal))
rs, err = c.Execute(ctx, partitionsQuery, minVal)
} else {
partitionsQuery := fmt.Sprintf(
`SELECT bucket, MIN(%[2]s) AS start, MAX(%[2]s) AS end
FROM (
SELECT NTILE(%[1]d) OVER (ORDER BY %[2]s) AS bucket, %[2]s FROM %[3]s
) AS subquery
ORDER BY start;`,
numPartitions,
quotedWatermarkColumn,
parsedWatermarkTable.MySQL(),
)
c.logger.Info("partitions query", slog.String("query", partitionsQuery), slog.Any("minVal", minVal))
rs, err = c.Execute(ctx, partitionsQuery, minVal, minVal)
} else {
partitionsQuery := fmt.Sprintf(
`WITH stats AS (
SELECT MIN(%[2]s) AS min_watermark,
1.0 * (MAX(%[2]s) - MIN(%[2]s)) / (%[1]d) AS range_size
FROM %[3]s
)
SELECT FLOOR((w.%[2]s - s.min_watermark) / s.range_size) AS bucket,
MIN(w.%[2]s) AS start, MAX(w.%[2]s) AS end
FROM %[3]s AS w
CROSS JOIN stats AS s
GROUP BY bucket
ORDER BY start`,
numPartitions,
quotedWatermarkColumn,
parsedWatermarkTable.MySQL(),
)
c.logger.Info("partitions query", slog.String("query", partitionsQuery))
rs, err = c.Execute(ctx, partitionsQuery)
}
if err != nil {
return nil, fmt.Errorf("failed to query for partitions: %w", err)
ORDER BY start;`,
numPartitions,
quotedWatermarkColumn,
parsedWatermarkTable.MySQL(),
)
c.logger.Info("partitions query", slog.String("query", partitionsQuery))
rs, err = c.Execute(ctx, partitionsQuery)
}
if err != nil {
return nil, fmt.Errorf("failed to query for partitions: %w", err)
}
case types.QValueKindTimestamp, types.QValueKindTimestampTZ:
if minVal != nil {
partitionsQuery := fmt.Sprintf(
`WITH stats AS (
SELECT MIN(%[2]s) AS min_watermark,
1.0 * (TIMESTAMPDIFF(MICROSECOND, MAX(%[2]s), MIN(%[2]s)) / (%[1]d)) AS range_size
FROM %[3]s WHERE %[2]s > ?
)
SELECT FLOOR(TIMESTAMPDIFF(MICROSECOND, w.%[2]s, s.min_watermark) / s.range_size) AS bucket,
MIN(w.%[2]s) AS start, MAX(w.%[2]s) AS end
FROM %[3]s AS w
CROSS JOIN stats AS s
WHERE w.%[2]s > ?
GROUP BY bucket
ORDER BY start;`,
numPartitions,
quotedWatermarkColumn,
parsedWatermarkTable.MySQL(),
)
c.logger.Info("partitions query", slog.String("query", partitionsQuery), slog.Any("minVal", minVal))
rs, err = c.Execute(ctx, partitionsQuery, minVal)
} else {
partitionsQuery := fmt.Sprintf(
`WITH stats AS (
SELECT MIN(%[2]s) AS min_watermark,
1.0 * (TIMESTAMPDIFF(MICROSECOND, MAX(%[2]s), MIN(%[2]s)) / (%[1]d)) AS range_size
FROM %[3]s
)
SELECT FLOOR(TIMESTAMPDIFF(MICROSECOND, w.%[2]s, s.min_watermark) / s.range_size) AS bucket,
MIN(w.%[2]s) AS start, MAX(w.%[2]s) AS end
FROM %[3]s AS w
CROSS JOIN stats AS s
GROUP BY bucket
ORDER BY start;`,
numPartitions,
quotedWatermarkColumn,
parsedWatermarkTable.MySQL(),
)
c.logger.Info("partitions query", slog.String("query", partitionsQuery))
rs, err = c.Execute(ctx, partitionsQuery)
}
if err != nil {
return nil, fmt.Errorf("failed to query for partitions: %w", err)
}
}

qk1, err := qkindFromMysql(rs.Fields[1])
if err != nil {
return nil, err
}
qk2, err := qkindFromMysql(rs.Fields[2])
if err != nil {
return nil, err
}
if qk1 != qk2 {
return nil, fmt.Errorf("low/high of partition range should be same type, got low:%s, high:%s", qk1, qk2)
}
partitionHelper := utils.NewPartitionHelper(c.logger)
for _, row := range rs.Values {
val1, err := QValueFromMysqlFieldValue(qk1, row[1])
val1, err := QValueFromMysqlFieldValue(watermarkQKind, row[1])
if err != nil {
return nil, err
}
val2, err := QValueFromMysqlFieldValue(qk2, row[2])
val2, err := QValueFromMysqlFieldValue(watermarkQKind, row[2])
if err != nil {
return nil, err
}
Expand All @@ -169,7 +244,7 @@ func (c *MySqlConnector) GetQRepPartitions(
func (c *MySqlConnector) PullQRepRecords(
ctx context.Context,
config *protos.QRepConfig,
last *protos.QRepPartition,
partition *protos.QRepPartition,
stream *model.QRecordStream,
) (int64, int64, error) {
tableSchema, err := c.getTableSchemaForTable(ctx, config.Env,
Expand Down Expand Up @@ -241,7 +316,7 @@ func (c *MySqlConnector) PullQRepRecords(
return nil
}

if last.FullTablePartition {
if partition.FullTablePartition {
// this is a full table partition, so just run the query
if err := c.ExecuteSelectStreaming(ctx, config.Query, &rs, onRow, onResult); err != nil {
return 0, 0, err
Expand All @@ -251,7 +326,7 @@ func (c *MySqlConnector) PullQRepRecords(
var rangeEnd string

// Depending on the type of the range, convert the range into the correct type
switch x := last.Range.Range.(type) {
switch x := partition.Range.Range.(type) {
case *protos.PartitionRange_IntRange:
rangeStart = strconv.FormatInt(x.IntRange.Start, 10)
rangeEnd = strconv.FormatInt(x.IntRange.End, 10)
Expand Down
Loading