Skip to content

Commit 869cb93

Browse files
thesyncimthesyncim
committed
perf(cdc): replay updates as primary-key upserts
1 parent 22b0007 commit 869cb93

3 files changed

Lines changed: 399 additions & 24 deletions

File tree

‎internal/cdc/applier.go‎

Lines changed: 311 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,7 @@ type targetColumn struct {
454454
oid uint32
455455
arrayOID uint32
456456
key bool
457+
primary bool
457458
identity string
458459
sourceIndex int
459460
generated bool
@@ -947,6 +948,12 @@ func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *R
947948
a.attidentity::text,
948949
a.attgenerated <> '',
949950
a.attnotnull,
951+
EXISTS (
952+
SELECT 1 FROM pg_catalog.pg_index primary_index
953+
WHERE primary_index.indrelid = c.oid
954+
AND primary_index.indisprimary
955+
AND a.attnum = ANY(primary_index.indkey)
956+
) AS primary_key,
950957
EXISTS (
951958
SELECT 1 FROM pg_catalog.pg_index conflict_index
952959
WHERE conflict_index.indrelid = c.oid
@@ -1017,7 +1024,7 @@ func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *R
10171024
var heapBytes, heapBlocksRead, heapBlocksHit int64
10181025
if err := rows.Scan(
10191026
&column.name, &column.oid, &column.arrayOID, &column.identity,
1020-
&column.generated, &column.notNull, &column.conflicting,
1027+
&column.generated, &column.notNull, &column.primary, &column.conflicting,
10211028
&selectiveUpdates, &setDMLSafe, &builtIn, &heapBytes,
10221029
&heapBlocksRead, &heapBlocksHit,
10231030
); err != nil {
@@ -2223,7 +2230,307 @@ func applyInsertArrayChunk(
22232230
})
22242231
}
22252232

2233+
// PostgreSQL logical replication supplies a complete new row for ordinary
2234+
// columns (apart from unchanged TOAST values). For those rows, use the same
2235+
// primary-key upsert shape as crdb-to-pg: PostgreSQL resolves the conflict
2236+
// through the exact primary key and performs the update in one operation. This
2237+
// avoids a compare-first target read and avoids UPDATE ... FROM plans whose
2238+
// join order can select an unrelated secondary index on very large tables.
2239+
func canPrimaryKeyUpsert(relation *targetRelation, change *Change) bool {
2240+
if relation == nil || change == nil || change.New == nil ||
2241+
len(*change.New) != len(relation.source.Columns) || len(relation.columns) == 0 ||
2242+
!relation.capabilities.keyedSetDML {
2243+
return false
2244+
}
2245+
primary := primaryKeyColumns(relation)
2246+
if len(primary) == 0 {
2247+
return false
2248+
}
2249+
for _, column := range relation.columns {
2250+
if (*change.New)[column.sourceIndex].Kind == DatumUnchangedToast {
2251+
return false
2252+
}
2253+
}
2254+
if change.Old == nil || len(*change.Old) != len(relation.source.Columns) {
2255+
return true
2256+
}
2257+
for _, column := range primary {
2258+
oldDatum := (*change.Old)[column.sourceIndex]
2259+
newDatum := (*change.New)[column.sourceIndex]
2260+
if oldDatum.Kind == DatumUnchangedToast || !tupleDatumEqual(oldDatum, newDatum) {
2261+
return false
2262+
}
2263+
}
2264+
return true
2265+
}
2266+
2267+
func tupleDatumEqual(left, right TupleDatum) bool {
2268+
return left.Kind == right.Kind && bytes.Equal(left.Data, right.Data)
2269+
}
2270+
2271+
func primaryKeyColumns(relation *targetRelation) []targetColumn {
2272+
columns := make([]targetColumn, 0, len(relation.columns))
2273+
for _, column := range relation.columns {
2274+
if column.primary {
2275+
columns = append(columns, column)
2276+
}
2277+
}
2278+
return columns
2279+
}
2280+
2281+
func primaryKeyTupleKey(relation *targetRelation, tuple *Tuple) (string, error) {
2282+
if err := validateTuple(relation, tuple, ChangeUpdate); err != nil {
2283+
return "", err
2284+
}
2285+
var key strings.Builder
2286+
for _, column := range primaryKeyColumns(relation) {
2287+
datum := (*tuple)[column.sourceIndex]
2288+
key.WriteByte(byte(datum.Kind))
2289+
fmt.Fprintf(&key, ":%d:", len(datum.Data))
2290+
key.Write(datum.Data)
2291+
}
2292+
return key.String(), nil
2293+
}
2294+
2295+
func applyPrimaryKeyUpsertChunk(
2296+
replay *applyPipeline,
2297+
relation *targetRelation,
2298+
changes []Change,
2299+
) error {
2300+
if len(changes) == 0 {
2301+
return nil
2302+
}
2303+
if applied, err := applyPrimaryKeyUpsertTextStage(replay, relation, changes); applied || err != nil {
2304+
return err
2305+
}
2306+
if applied, err := applyPrimaryKeyUpsertArrayChunk(replay, relation, changes); applied || err != nil {
2307+
return err
2308+
}
2309+
chunkRows := insertChunkRows(len(relation.columns))
2310+
for start := 0; start < len(changes); start += chunkRows {
2311+
end := min(start+chunkRows, len(changes))
2312+
if err := applyPrimaryKeyUpsertValueChunk(replay, relation, changes[start:end]); err != nil {
2313+
return err
2314+
}
2315+
}
2316+
return nil
2317+
}
2318+
2319+
func applyPrimaryKeyUpsertTextStage(
2320+
replay *applyPipeline,
2321+
relation *targetRelation,
2322+
changes []Change,
2323+
) (bool, error) {
2324+
values := make([]TupleDatum, 0, len(changes)*len(relation.columns))
2325+
for row := range changes {
2326+
if err := validateTuple(relation, changes[row].New, ChangeUpdate); err != nil {
2327+
return true, err
2328+
}
2329+
for _, column := range relation.columns {
2330+
values = append(values, (*changes[row].New)[column.sourceIndex])
2331+
}
2332+
}
2333+
stage, applied, err := replay.loadTextCopyStage(
2334+
relation, ChangeUpdate, relation.columns, values, len(changes),
2335+
)
2336+
if err != nil || !applied {
2337+
return applied, err
2338+
}
2339+
var sql strings.Builder
2340+
writePrimaryKeyUpsertPrefix(&sql, relation)
2341+
sql.WriteString(" SELECT ")
2342+
for i := range relation.columns {
2343+
if i != 0 {
2344+
sql.WriteByte(',')
2345+
}
2346+
fmt.Fprintf(&sql, "column_%d", i)
2347+
}
2348+
sql.WriteString(" FROM ")
2349+
sql.WriteString(stage)
2350+
sql.WriteString(" ORDER BY ordinal")
2351+
appendPrimaryKeyConflictClause(&sql, relation)
2352+
return true, replay.queue(sql.String(), nil, applyExpectation{
2353+
relation: relation, kind: ChangeUpdate,
2354+
description: "staged primary-key upsert into " + relation.quoted,
2355+
expectedRows: int64(len(changes)),
2356+
})
2357+
}
2358+
2359+
func applyPrimaryKeyUpsertArrayChunk(
2360+
replay *applyPipeline,
2361+
relation *targetRelation,
2362+
changes []Change,
2363+
) (bool, error) {
2364+
params := make([]rawParam, 0, len(relation.columns))
2365+
for _, column := range relation.columns {
2366+
datums := make([]TupleDatum, len(changes))
2367+
for row := range changes {
2368+
if err := validateTuple(relation, changes[row].New, ChangeUpdate); err != nil {
2369+
return true, err
2370+
}
2371+
datums[row] = (*changes[row].New)[column.sourceIndex]
2372+
}
2373+
param, supported, err := arrayParamForColumn(relation, column, datums, ChangeUpdate)
2374+
if err != nil || !supported {
2375+
return supported, err
2376+
}
2377+
params = append(params, param)
2378+
}
2379+
var sql strings.Builder
2380+
writePrimaryKeyUpsertPrefix(&sql, relation)
2381+
sql.WriteString(" SELECT ")
2382+
for i := range relation.columns {
2383+
if i != 0 {
2384+
sql.WriteByte(',')
2385+
}
2386+
fmt.Fprintf(&sql, "pgmigrate_batch.column_%d", i)
2387+
}
2388+
sql.WriteString(" FROM unnest(")
2389+
for i := range params {
2390+
if i != 0 {
2391+
sql.WriteByte(',')
2392+
}
2393+
fmt.Fprintf(&sql, "$%d", i+1)
2394+
}
2395+
sql.WriteString(") AS pgmigrate_batch(")
2396+
for i := range relation.columns {
2397+
if i != 0 {
2398+
sql.WriteByte(',')
2399+
}
2400+
fmt.Fprintf(&sql, "column_%d", i)
2401+
}
2402+
sql.WriteString(") WHERE true")
2403+
appendPrimaryKeyConflictClause(&sql, relation)
2404+
return true, replay.queue(sql.String(), params, applyExpectation{
2405+
relation: relation, kind: ChangeUpdate,
2406+
description: "array primary-key upsert into " + relation.quoted,
2407+
expectedRows: int64(len(changes)),
2408+
})
2409+
}
2410+
2411+
func applyPrimaryKeyUpsertValueChunk(
2412+
replay *applyPipeline,
2413+
relation *targetRelation,
2414+
changes []Change,
2415+
) error {
2416+
var sql strings.Builder
2417+
writePrimaryKeyUpsertPrefix(&sql, relation)
2418+
sql.WriteString(" VALUES ")
2419+
params := make([]rawParam, 0, len(changes)*len(relation.columns))
2420+
for row := range changes {
2421+
if err := validateTuple(relation, changes[row].New, ChangeUpdate); err != nil {
2422+
return err
2423+
}
2424+
if row != 0 {
2425+
sql.WriteByte(',')
2426+
}
2427+
sql.WriteByte('(')
2428+
for columnIndex, column := range relation.columns {
2429+
if columnIndex != 0 {
2430+
sql.WriteByte(',')
2431+
}
2432+
param, err := datumParamForColumn(
2433+
relation, column, (*changes[row].New)[column.sourceIndex], ChangeUpdate,
2434+
)
2435+
if err != nil {
2436+
return err
2437+
}
2438+
params = append(params, param)
2439+
fmt.Fprintf(&sql, "$%d", len(params))
2440+
}
2441+
sql.WriteByte(')')
2442+
}
2443+
appendPrimaryKeyConflictClause(&sql, relation)
2444+
return replay.queue(sql.String(), params, applyExpectation{
2445+
relation: relation, kind: ChangeUpdate,
2446+
description: "primary-key upsert into " + relation.quoted,
2447+
expectedRows: int64(len(changes)),
2448+
})
2449+
}
2450+
2451+
func writePrimaryKeyUpsertPrefix(sql *strings.Builder, relation *targetRelation) {
2452+
sql.WriteString("INSERT INTO ")
2453+
sql.WriteString(relation.quoted)
2454+
sql.WriteString(" (")
2455+
for i, column := range relation.columns {
2456+
if i != 0 {
2457+
sql.WriteByte(',')
2458+
}
2459+
sql.WriteString(column.quoted)
2460+
}
2461+
sql.WriteByte(')')
2462+
if relation.overrideIdentity {
2463+
sql.WriteString(" OVERRIDING SYSTEM VALUE")
2464+
}
2465+
}
2466+
2467+
func appendPrimaryKeyConflictClause(sql *strings.Builder, relation *targetRelation) {
2468+
primary := primaryKeyColumns(relation)
2469+
sql.WriteString(" ON CONFLICT (")
2470+
for i, column := range primary {
2471+
if i != 0 {
2472+
sql.WriteByte(',')
2473+
}
2474+
sql.WriteString(column.quoted)
2475+
}
2476+
sql.WriteString(") DO UPDATE SET ")
2477+
assignments := 0
2478+
for _, column := range relation.columns {
2479+
if column.primary {
2480+
continue
2481+
}
2482+
if assignments != 0 {
2483+
sql.WriteByte(',')
2484+
}
2485+
sql.WriteString(column.quoted)
2486+
sql.WriteString("=EXCLUDED.")
2487+
sql.WriteString(column.quoted)
2488+
assignments++
2489+
}
2490+
if assignments == 0 {
2491+
sql.WriteString(primary[0].quoted)
2492+
sql.WriteString("=EXCLUDED.")
2493+
sql.WriteString(primary[0].quoted)
2494+
}
2495+
}
2496+
22262497
func applyUpdates(replay *applyPipeline, relation *targetRelation, changes []Change) error {
2498+
for start := 0; start < len(changes); {
2499+
if !canPrimaryKeyUpsert(relation, &changes[start]) {
2500+
end := start + 1
2501+
for end < len(changes) && !canPrimaryKeyUpsert(relation, &changes[end]) {
2502+
end++
2503+
}
2504+
if err := applyLegacyUpdates(replay, relation, changes[start:end]); err != nil {
2505+
return err
2506+
}
2507+
start = end
2508+
continue
2509+
}
2510+
2511+
seen := make(map[string]struct{})
2512+
end := start
2513+
for end < len(changes) && end-start < applyArrayChunkRows &&
2514+
canPrimaryKeyUpsert(relation, &changes[end]) {
2515+
key, err := primaryKeyTupleKey(relation, changes[end].New)
2516+
if err != nil {
2517+
return err
2518+
}
2519+
if _, duplicate := seen[key]; duplicate {
2520+
break
2521+
}
2522+
seen[key] = struct{}{}
2523+
end++
2524+
}
2525+
if err := applyPrimaryKeyUpsertChunk(replay, relation, changes[start:end]); err != nil {
2526+
return err
2527+
}
2528+
start = end
2529+
}
2530+
return nil
2531+
}
2532+
2533+
func applyLegacyUpdates(replay *applyPipeline, relation *targetRelation, changes []Change) error {
22272534
identityColumns := batchUpdateIdentityColumns(relation)
22282535
if len(changes) < 2 || len(identityColumns) == 0 || len(relation.columns) == 0 {
22292536
for i := range changes {
@@ -3071,7 +3378,7 @@ func applyUpdateTextStage(
30713378
sql.WriteString(" FROM ")
30723379
sql.WriteString(stage)
30733380
sql.WriteString(" AS pgmigrate_batch WHERE ")
3074-
if len(identityColumns) > 1 {
3381+
if len(identityColumns) > 1 && useSelectiveBitmap(relation) {
30753382
writeCompositeIdentityCTIDPredicate(
30763383
&sql, relation, identityColumns, "column_", len(setColumns),
30773384
)
@@ -3153,7 +3460,7 @@ func applyUpdateValueChunk(
31533460
fmt.Fprintf(&sql, ",identity_%d", i)
31543461
}
31553462
sql.WriteString(") WHERE ")
3156-
if len(identityColumns) > 1 {
3463+
if len(identityColumns) > 1 && useSelectiveBitmap(relation) {
31573464
writeCompositeIdentityCTIDPredicate(
31583465
&sql, relation, identityColumns, "identity_", 0,
31593466
)
@@ -3269,7 +3576,7 @@ func applyUpdateArrayChunk(
32693576
sql.WriteByte(',')
32703577
}
32713578
sql.WriteString("ordinal) WHERE ")
3272-
if len(identityColumns) > 1 {
3579+
if len(identityColumns) > 1 && useSelectiveBitmap(relation) {
32733580
writeCompositeIdentityCTIDPredicate(
32743581
&sql, relation, identityColumns, "identity_", 0,
32753582
)

0 commit comments

Comments
 (0)