Skip to content

Commit a1debd7

Browse files
jgao54jgao54
committed
introduce min/max-id based parallel partitioning for mongodb
1 parent e9c5d3f commit a1debd7

2 files changed

Lines changed: 245 additions & 67 deletions

File tree

‎flow/connectors/mongo/qrep.go‎

Lines changed: 40 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ import (
88
"math"
99
"time"
1010

11-
"github.com/google/uuid"
1211
"go.mongodb.org/mongo-driver/v2/bson"
12+
"go.mongodb.org/mongo-driver/v2/mongo"
1313
"go.mongodb.org/mongo-driver/v2/mongo/options"
1414

1515
"github.com/PeerDB-io/peerdb/flow/connectors/utils"
@@ -38,6 +38,10 @@ func (c *MongoConnector) GetQRepPartitions(
3838
return fullTablePartition, nil
3939
}
4040

41+
if config.WatermarkColumn != DefaultDocumentKeyColumnName {
42+
return nil, fmt.Errorf("only %s is currently supported as watermark column for MongoDB connector", DefaultDocumentKeyColumnName)
43+
}
44+
4145
if config.NumRowsPerPartition <= 0 {
4246
return nil, fmt.Errorf("num rows per partition must be greater than 0")
4347
} else if last != nil && last.Range != nil {
@@ -58,9 +62,6 @@ func (c *MongoConnector) GetQRepPartitions(
5862
if err != nil {
5963
return nil, fmt.Errorf("failed to count documents in collection %s: %w", parseWatermarkTable.Table, err)
6064
}
61-
if totalRows == 0 {
62-
return []*protos.QRepPartition{}, nil
63-
}
6465

6566
// Calculate the number of partitions
6667
adjustedPartitions := shared.AdjustNumPartitions(totalRows, numRowsPerPartition)
@@ -70,78 +71,50 @@ func (c *MongoConnector) GetQRepPartitions(
7071
slog.Int64("adjustedNumPartitions", adjustedPartitions.AdjustedNumPartitions),
7172
slog.Int64("adjustedNumRowsPerPartition", adjustedPartitions.AdjustedNumRowsPerPartition))
7273

73-
// no need to bother with bucketAuto if we have only one partition
74-
if adjustedPartitions.AdjustedNumPartitions == 1 {
75-
return fullTablePartition, nil
76-
}
77-
78-
if config.WatermarkColumn != DefaultDocumentKeyColumnName {
79-
return nil, fmt.Errorf("only %s is currently supported as watermark column for MongoDB connector", DefaultDocumentKeyColumnName)
80-
}
74+
return c.minMaxPartitions(ctx, collection, adjustedPartitions.AdjustedNumPartitions)
75+
}
8176

82-
// Use bucketAuto to create partitions based on _id field
83-
bucketAutoPipeline := []bson.D{
84-
{
85-
{Key: "$bucketAuto", Value: bson.D{
86-
{Key: "groupBy", Value: "$" + config.WatermarkColumn},
87-
{Key: "buckets", Value: adjustedPartitions.AdjustedNumPartitions},
88-
}},
89-
},
77+
func (c *MongoConnector) GetDefaultPartitionKeyForTables(
78+
ctx context.Context,
79+
input *protos.GetDefaultPartitionKeyForTablesInput,
80+
) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
81+
mapping := make(map[string]string, len(input.TableMappings))
82+
for _, tm := range input.TableMappings {
83+
parsedTable, err := common.ParseTableIdentifier(tm.SourceTableIdentifier)
84+
if err != nil {
85+
c.logger.Warn("[mongo] failed to parse table identifier, skipping parallel load",
86+
slog.String("table", tm.SourceTableIdentifier),
87+
slog.String("error", err.Error()))
88+
continue
89+
}
90+
collection := c.client.Database(parsedTable.Namespace).Collection(parsedTable.Table)
91+
if hasObjectIDAsKey(ctx, collection) {
92+
mapping[tm.SourceTableIdentifier] = DefaultDocumentKeyColumnName
93+
} else {
94+
c.logger.Info("[mongo] _id is not ObjectID type, falling back to full table partition",
95+
slog.String("table", tm.SourceTableIdentifier))
96+
}
9097
}
98+
return &protos.GetDefaultPartitionKeyForTablesOutput{
99+
TableDefaultPartitionKeyMapping: mapping,
100+
}, nil
101+
}
91102

92-
cursor, err := collection.Aggregate(ctx, bucketAutoPipeline)
103+
func hasObjectIDAsKey(ctx context.Context, collection *mongo.Collection) bool {
104+
cursor, err := collection.Find(ctx, bson.D{}, options.Find().SetLimit(1).SetProjection(bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}))
93105
if err != nil {
94-
return nil, fmt.Errorf("failed to aggregate for bucket partitions: %w", err)
106+
return false
95107
}
96108
defer cursor.Close(ctx)
97109

98-
partitions := make([]*protos.QRepPartition, 0, adjustedPartitions.AdjustedNumPartitions)
99-
for cursor.Next(ctx) {
100-
var bucket struct {
101-
ID struct {
102-
Min bson.ObjectID `bson:"min"`
103-
Max bson.ObjectID `bson:"max"`
104-
} `bson:"_id"`
105-
}
106-
if err := bson.Unmarshal(cursor.Current, &bucket); err != nil {
107-
return nil, fmt.Errorf("failed to unmarshal bucket: %w", err)
108-
}
109-
110-
partitions = append(partitions, &protos.QRepPartition{
111-
PartitionId: uuid.NewString(),
112-
Range: &protos.PartitionRange{
113-
Range: &protos.PartitionRange_ObjectIdRange{
114-
ObjectIdRange: &protos.ObjectIdPartitionRange{
115-
Start: bucket.ID.Min.Hex(),
116-
End: bucket.ID.Max.Hex(),
117-
},
118-
},
119-
},
120-
FullTablePartition: false,
121-
})
110+
if !cursor.Next(ctx) {
111+
return false
122112
}
123-
if err := cursor.Err(); err != nil {
124-
if errors.Is(err, context.Canceled) {
125-
c.logger.Warn("context canceled while performing bucketAuto aggregation",
126-
slog.String("watermark_table", config.WatermarkTable))
127-
} else {
128-
c.logger.Error("error while performing bucketAuto aggregation",
129-
slog.String("watermark_table", config.WatermarkTable),
130-
slog.String("error", err.Error()))
131-
}
132-
return nil, fmt.Errorf("cursor error during bucketAuto aggregation: %w", err)
113+
var oid bson.ObjectID
114+
if err := bson.Unmarshal(cursor.Current, &oid); err != nil {
115+
return false
133116
}
134-
135-
return partitions, nil
136-
}
137-
138-
func (c *MongoConnector) GetDefaultPartitionKeyForTables(
139-
ctx context.Context,
140-
input *protos.GetDefaultPartitionKeyForTablesInput,
141-
) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
142-
return &protos.GetDefaultPartitionKeyForTablesOutput{
143-
TableDefaultPartitionKeyMapping: nil,
144-
}, nil
117+
return !oid.IsZero()
145118
}
146119

147120
func (c *MongoConnector) PullQRepRecords(

‎flow/connectors/mongo/qrep_partition.go‎

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
package connmongo
2+
3+
import (
4+
"context"
5+
"encoding/binary"
6+
"errors"
7+
"fmt"
8+
"log/slog"
9+
"time"
10+
11+
"github.com/google/uuid"
12+
"go.mongodb.org/mongo-driver/v2/bson"
13+
"go.mongodb.org/mongo-driver/v2/mongo"
14+
"go.mongodb.org/mongo-driver/v2/mongo/options"
15+
16+
"github.com/PeerDB-io/peerdb/flow/connectors/utils"
17+
"github.com/PeerDB-io/peerdb/flow/generated/protos"
18+
"github.com/PeerDB-io/peerdb/flow/shared"
19+
)
20+
21+
// objectIDMinForTimestamp creates an ObjectID with the given timestamp and
22+
// all trailing bytes set to 0x00, producing the smallest possible ObjectID
23+
// for that second.
24+
func objectIDMinForTimestamp(t time.Time) bson.ObjectID {
25+
var oid [12]byte
26+
binary.BigEndian.PutUint32(oid[0:4], uint32(t.Unix()))
27+
return oid
28+
}
29+
30+
// objectIDMaxForTimestamp creates an ObjectID with the given timestamp and
31+
// all trailing bytes set to 0xFF, producing the largest possible ObjectID
32+
// for that second.
33+
func objectIDMaxForTimestamp(t time.Time) bson.ObjectID {
34+
var oid [12]byte
35+
binary.BigEndian.PutUint32(oid[0:4], uint32(t.Unix()))
36+
for i := 4; i < 12; i++ {
37+
oid[i] = 0xFF
38+
}
39+
return oid
40+
}
41+
42+
// minMaxPartitions creates partitions by querying only the min and max _id
43+
// from the collection (leveraging the default _id index), then uniformly
44+
// dividing the timestamp range encoded in the ObjectIDs.
45+
func (c *MongoConnector) minMaxPartitions(
46+
ctx context.Context,
47+
collection *mongo.Collection,
48+
numPartitions int64,
49+
) ([]*protos.QRepPartition, error) {
50+
fullTablePartition := []*protos.QRepPartition{{
51+
PartitionId: utils.FullTablePartitionID,
52+
Range: nil,
53+
FullTablePartition: true,
54+
}}
55+
if numPartitions == 1 {
56+
c.logger.Info("[mongo] using full table partition for single partition")
57+
return fullTablePartition, nil
58+
}
59+
60+
minID, err := findBoundaryObjectID(ctx, collection, 1)
61+
if err != nil {
62+
return nil, fmt.Errorf("failed to find min _id: %w", err)
63+
}
64+
maxID, err := findBoundaryObjectID(ctx, collection, -1)
65+
if err != nil {
66+
return nil, fmt.Errorf("failed to find max _id: %w", err)
67+
}
68+
69+
tsMin := minID.Timestamp()
70+
tsMax := maxID.Timestamp()
71+
tsRange := tsMax.Unix() - tsMin.Unix()
72+
if tsRange <= 0 {
73+
c.logger.Info("[mongo] min/max timestamps are equal, returning single partition")
74+
return fullTablePartition, nil
75+
}
76+
77+
c.logger.Info("[mongo] using min/max ObjectID timestamp partitioning",
78+
slog.Time("tsMin", tsMin),
79+
slog.Time("tsMax", tsMax),
80+
slog.Int64("numPartitions", numPartitions))
81+
82+
secondsPerPartition := shared.DivCeil(tsRange, numPartitions)
83+
partitions := make([]*protos.QRepPartition, 0, numPartitions)
84+
for i := range numPartitions {
85+
var start, end bson.ObjectID
86+
if i == 0 {
87+
start = minID
88+
} else {
89+
start = objectIDMinForTimestamp(time.Unix(tsMin.Unix()+secondsPerPartition*i, 0))
90+
}
91+
if i == numPartitions-1 {
92+
end = maxID
93+
} else {
94+
end = objectIDMaxForTimestamp(time.Unix(tsMin.Unix()+secondsPerPartition*(i+1)-1, 0))
95+
}
96+
partitions = append(partitions, &protos.QRepPartition{
97+
PartitionId: uuid.NewString(),
98+
Range: &protos.PartitionRange{
99+
Range: &protos.PartitionRange_ObjectIdRange{
100+
ObjectIdRange: &protos.ObjectIdPartitionRange{
101+
Start: start.Hex(),
102+
End: end.Hex(),
103+
},
104+
},
105+
},
106+
FullTablePartition: false,
107+
})
108+
}
109+
110+
return partitions, nil
111+
}
112+
113+
// findBoundaryObjectID returns the min (direction=1) or max (direction=-1) _id
114+
// in the collection. Because _id is always indexed, this is an efficient O(log n) operation.
115+
func findBoundaryObjectID(
116+
ctx context.Context,
117+
collection *mongo.Collection,
118+
direction int,
119+
) (bson.ObjectID, error) {
120+
findCmd := bson.D{
121+
{Key: "find", Value: collection.Name()},
122+
{Key: "sort", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: direction}}},
123+
{Key: "limit", Value: 1},
124+
{Key: "projection", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}},
125+
}
126+
127+
cursor, err := collection.Database().RunCommandCursor(ctx, findCmd, options.RunCmd())
128+
if err != nil {
129+
return bson.NilObjectID, fmt.Errorf("failed to run find command: %w", err)
130+
}
131+
defer cursor.Close(ctx)
132+
133+
if !cursor.Next(ctx) {
134+
return bson.NilObjectID, fmt.Errorf("collection is empty")
135+
}
136+
var doc struct {
137+
ID bson.ObjectID `bson:"_id"`
138+
}
139+
if err := bson.Unmarshal(cursor.Current, &doc); err != nil {
140+
return bson.NilObjectID, fmt.Errorf("failed to unmarshal boundary document: %w", err)
141+
}
142+
return doc.ID, nil
143+
}
144+
145+
// bucketAutoPartitions uses MongoDB's $bucketAuto aggregation to create partitions.
146+
// This performs a full collection scan, so it is significantly slower on large collections
147+
// but produces well-balanced partitions.
148+
func (c *MongoConnector) bucketAutoPartitions(
149+
ctx context.Context,
150+
collection *mongo.Collection,
151+
watermarkColumn string,
152+
numPartitions int64,
153+
) ([]*protos.QRepPartition, error) {
154+
bucketAutoPipeline := []bson.D{
155+
{
156+
{Key: "$bucketAuto", Value: bson.D{
157+
{Key: "groupBy", Value: "$" + watermarkColumn},
158+
{Key: "buckets", Value: numPartitions},
159+
}},
160+
},
161+
}
162+
163+
cursor, err := collection.Aggregate(ctx, bucketAutoPipeline)
164+
if err != nil {
165+
return nil, fmt.Errorf("failed to aggregate for bucket partitions: %w", err)
166+
}
167+
defer cursor.Close(ctx)
168+
169+
partitions := make([]*protos.QRepPartition, 0, numPartitions)
170+
for cursor.Next(ctx) {
171+
var bucket struct {
172+
ID struct {
173+
Min bson.ObjectID `bson:"min"`
174+
Max bson.ObjectID `bson:"max"`
175+
} `bson:"_id"`
176+
}
177+
if err := bson.Unmarshal(cursor.Current, &bucket); err != nil {
178+
return nil, fmt.Errorf("failed to unmarshal bucket: %w", err)
179+
}
180+
181+
partitions = append(partitions, &protos.QRepPartition{
182+
PartitionId: uuid.NewString(),
183+
Range: &protos.PartitionRange{
184+
Range: &protos.PartitionRange_ObjectIdRange{
185+
ObjectIdRange: &protos.ObjectIdPartitionRange{
186+
Start: bucket.ID.Min.Hex(),
187+
End: bucket.ID.Max.Hex(),
188+
},
189+
},
190+
},
191+
FullTablePartition: false,
192+
})
193+
}
194+
if err := cursor.Err(); err != nil {
195+
if errors.Is(err, context.Canceled) {
196+
c.logger.Warn("context canceled while performing bucketAuto aggregation")
197+
} else {
198+
c.logger.Error("error while performing bucketAuto aggregation",
199+
slog.String("error", err.Error()))
200+
}
201+
return nil, fmt.Errorf("cursor error during bucketAuto aggregation: %w", err)
202+
}
203+
204+
return partitions, nil
205+
}

0 commit comments

Comments
 (0)