mongoDB parallel snapshot follow-up (#4531)

Follow-ups for MongoDB parallel snapshotting #4487 to introduce NumericPartitions; this avoids missing data during initial snapshots when there are interior fractional _ids.

A few other smaller clean-up, with inline comments.

Author: jgao54

flow/connectors/mongo/qrep.go

@@ -226,19 +226,18 @@ func toRangeFilter(watermarkColumn string, partitionRange *protos.PartitionRange
 				bson.E{Key: "$lte", Value: endObjectID},
 			}},
 		}, nil
-	case *protos.PartitionRange_IntRange:
-		// Numeric _id: inclusive range. PartitionHelper builds non-overlapping
-		// [start, end] integer ranges. Mongo compares int32/int64 numerically, so
-		// int64 bounds match int32-typed _ids.
+	case *protos.PartitionRange_NumericRange:
+		endOp := "$lt"
+		if r.NumericRange.EndInclusive {
+			endOp = "$lte"
+		}
 		return bson.D{
 			bson.E{Key: watermarkColumn, Value: bson.D{
-				bson.E{Key: "$gte", Value: r.IntRange.Start},
-				bson.E{Key: "$lte", Value: r.IntRange.End},
+				bson.E{Key: "$gte", Value: r.NumericRange.Start},
+				bson.E{Key: endOp, Value: r.NumericRange.End},
 			}},
 		}, nil
 	case *protos.PartitionRange_StringRange:
-		// String _id: half-open [start, end) for all but the last partition;
-		// the last partition is closed [start, end] (EndInclusive == true).
 		endOp := "$lt"
 		if r.StringRange.EndInclusive {
 			endOp = "$lte"

flow/connectors/mongo/qrep_partition.go

@@ -27,12 +27,10 @@ const (
 // buildPartitions reads the collection's min and max _id (cheap, leverages the
 // default _id index) and dispatches to a type-specific partitioner:
 //   - ObjectID  -> uniform division of the 12-byte ObjectID keyspace
-//   - Int32/64  -> uniform division of the numeric range (reuses PartitionHelper)
+//   - Int32/64  -> uniform division of the numeric range
 //   - String    -> sampling-based quantile boundaries
 //
-// Mixed-type _id (min and max differ in type), Double, or any other type fall back
-// to a full-table partition. The change stream resume token is captured before this
-// runs, so any writes during partitioning are replayed by CDC.
+// Mixed-type _id, Double, or any other type fall back to a full-table partition.
 func (c *MongoConnector) buildPartitions(
 	ctx context.Context,
 	collection *mongo.Collection,
@@ -127,15 +125,10 @@ func (c *MongoConnector) objectIDPartitions(
 	return partitions, nil
 }
 
-// numericPartitions divides a numeric (int32/int64) _id range uniformly, reusing
-// the shared PartitionHelper which builds non-overlapping IntPartitionRange
-// partitions with tested +1 boundary semantics covering [min, max] inclusive.
-//
-// Limitation: IntPartitionRange uses inclusive-inclusive [start, end] boundaries,
-// so documents with double _id values that fall between two adjacent integer
-// boundaries (e.g. _id=2.5 between [1,2] and [3,4]) would be missed. This is an
-// edge case for collections with mixed int/double _ids; a follow-up will introduce
-// a NumericPartitionRange with half-open [start, end) semantics to address it.
+// numericPartitions divides a numeric (int32/int64) _id range uniformly. MongoDB
+// sorts all numeric types by value, so integer min/max _ids cannot rule out
+// fractional (double/decimal) _ids in-between. Therefore we use [start, end) to
+// ensure contiguous boundaries.
 func (c *MongoConnector) numericPartitions(
 	minVal int64,
 	maxVal int64,
@@ -146,16 +139,34 @@ func (c *MongoConnector) numericPartitions(
 		return utils.FullTablePartition(), nil
 	}
 
-	c.logger.Info("[mongo] using numeric _id partitioning",
-		slog.Int64("minID", minVal),
-		slog.Int64("maxID", maxVal),
-		slog.Int64("numPartitions", numPartitions))
+	ranges := computeNumericRanges(minVal, maxVal, numPartitions)
+	c.logger.Info("[mongo] using numeric _id partitioning", slog.Int("numPartitions", len(ranges)))
+	partitions := make([]*protos.QRepPartition, 0, len(ranges))
+	for i, rng := range ranges {
+		partitions = append(partitions, utils.CreateNumericPartition(rng[0], rng[1], i == len(ranges)-1))
+	}
+	return partitions, nil
+}
 
-	helper := utils.NewPartitionHelper(c.logger)
-	if err := helper.AddPartitionsWithRange(minVal, maxVal, numPartitions); err != nil {
-		return nil, fmt.Errorf("failed to build numeric partitions: %w", err)
+func computeNumericRanges(minVal int64, maxVal int64, numPartitions int64) [][2]int64 {
+	// span arithmetic is done in uint64 to avoid overflow
+	span := uint64(maxVal) - uint64(minVal)
+	step := span / uint64(numPartitions)
+	if span%uint64(numPartitions) != 0 {
+		step++
+	}
+
+	ranges := make([][2]int64, 0, numPartitions)
+	start := minVal
+	for {
+		if remaining := uint64(maxVal) - uint64(start); remaining <= step {
+			ranges = append(ranges, [2]int64{start, maxVal})
+			return ranges
+		}
+		end := int64(uint64(start) + step)
+		ranges = append(ranges, [2]int64{start, end})
+		start = end
 	}
-	return helper.GetPartitions(), nil
 }
 
 // stringPartitions builds partitions for a string _id collection. Because string
@@ -187,26 +198,12 @@ func (c *MongoConnector) stringPartitions(
 	}
 
 	c.logger.Info("[mongo] using sampled string _id partitioning",
-		slog.String("minID", minVal),
-		slog.String("maxID", maxVal),
 		slog.Int("numPartitions", len(ranges)),
 		slog.Int("sampleCount", len(samples)))
 
 	partitions := make([]*protos.QRepPartition, 0, len(ranges))
 	for i, rng := range ranges {
-		partitions = append(partitions, &protos.QRepPartition{
-			PartitionId: uuid.NewString(),
-			Range: &protos.PartitionRange{
-				Range: &protos.PartitionRange_StringRange{
-					StringRange: &protos.StringPartitionRange{
-						Start:        rng[0],
-						End:          rng[1],
-						EndInclusive: i == len(ranges)-1,
-					},
-				},
-			},
-			FullTablePartition: false,
-		})
+		partitions = append(partitions, utils.CreateStringPartition(rng[0], rng[1], i == len(ranges)-1))
 	}
 	return partitions, nil
 }
@@ -229,8 +226,8 @@ func (c *MongoConnector) sampleStringIDs(
 		{Key: "aggregate", Value: collection.Name()},
 		{Key: "pipeline", Value: bson.A{
 			bson.D{{Key: "$sample", Value: bson.D{{Key: "size", Value: int32(sampleSize)}}}},
-			bson.D{{Key: "$sort", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}}},
 			bson.D{{Key: "$project", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}}},
+			bson.D{{Key: "$sort", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}}},
 		}},
 		{Key: "cursor", Value: bson.D{}},
 	}

flow/connectors/mongo/qrep_partition_test.go

@@ -1,11 +1,13 @@
 package connmongo
 
 import (
+	"math"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"go.mongodb.org/mongo-driver/v2/bson"
 
+	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
 )
 
@@ -117,32 +119,76 @@ func TestComputeStringBoundaries(t *testing.T) {
 	})
 }
 
-func intPartition(start, end int64) *protos.PartitionRange {
-	return &protos.PartitionRange{Range: &protos.PartitionRange_IntRange{
-		IntRange: &protos.IntPartitionRange{Start: start, End: end},
-	}}
+func assertContiguousNumeric(t *testing.T, ranges [][2]int64, minVal, maxVal int64) {
+	t.Helper()
+	require.NotEmpty(t, ranges)
+	require.Equal(t, minVal, ranges[0][0])
+	require.Equal(t, maxVal, ranges[len(ranges)-1][1])
+	for i := range ranges {
+		require.Less(t, ranges[i][0], ranges[i][1])
+		if i+1 < len(ranges) {
+			require.Equal(t, ranges[i][1], ranges[i+1][0])
+		}
+	}
 }
 
-func stringPartition(start, end string, endInclusive bool) *protos.PartitionRange {
-	return &protos.PartitionRange{Range: &protos.PartitionRange_StringRange{
-		StringRange: &protos.StringPartitionRange{Start: start, End: end, EndInclusive: endInclusive},
-	}}
+func TestComputeNumericRanges(t *testing.T) {
+	t.Run("even division", func(t *testing.T) {
+		ranges := computeNumericRanges(1, 100, 10)
+		require.Len(t, ranges, 10)
+		assertContiguousNumeric(t, ranges, 1, 100)
+	})
+
+	t.Run("span smaller than partition count", func(t *testing.T) {
+		ranges := computeNumericRanges(0, 5, 10)
+		require.Len(t, ranges, 5)
+		assertContiguousNumeric(t, ranges, 0, 5)
+	})
+
+	t.Run("two adjacent values", func(t *testing.T) {
+		ranges := computeNumericRanges(7, 8, 4)
+		require.Len(t, ranges, 1)
+		assertContiguousNumeric(t, ranges, 7, 8)
+	})
+
+	t.Run("span does not overflow", func(t *testing.T) {
+		ranges := computeNumericRanges(math.MinInt64+5, math.MaxInt64-5, 10)
+		require.Len(t, ranges, 10)
+		assertContiguousNumeric(t, ranges, math.MinInt64+5, math.MaxInt64-5)
+	})
+
+	t.Run("full int64 domain does not overflow", func(t *testing.T) {
+		ranges := computeNumericRanges(math.MinInt64, math.MaxInt64, 7)
+		require.Len(t, ranges, 7)
+		assertContiguousNumeric(t, ranges, math.MinInt64, math.MaxInt64)
+	})
 }
 
 func TestToRangeFilter(t *testing.T) {
-	t.Run("int range is inclusive on both ends", func(t *testing.T) {
-		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, intPartition(10, 20))
+	t.Run("numeric range is half-open", func(t *testing.T) {
+		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, utils.CreateNumericPartition(10, 20, false).Range)
 		require.NoError(t, err)
 		require.Equal(t, bson.D{
 			bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{
 				bson.E{Key: "$gte", Value: int64(10)},
-				bson.E{Key: "$lte", Value: int64(20)},
+				bson.E{Key: "$lt", Value: int64(20)},
+			}},
+		}, filter)
+	})
+
+	t.Run("last numeric range is closed", func(t *testing.T) {
+		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, utils.CreateNumericPartition(20, 30, true).Range)
+		require.NoError(t, err)
+		require.Equal(t, bson.D{
+			bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{
+				bson.E{Key: "$gte", Value: int64(20)},
+				bson.E{Key: "$lte", Value: int64(30)},
 			}},
 		}, filter)
 	})
 
 	t.Run("string range is half-open", func(t *testing.T) {
-		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.a", "com.m", false))
+		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, utils.CreateStringPartition("com.a", "com.m", false).Range)
 		require.NoError(t, err)
 		require.Equal(t, bson.D{
 			bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{
@@ -153,7 +199,7 @@ func TestToRangeFilter(t *testing.T) {
 	})
 
 	t.Run("last string range is closed", func(t *testing.T) {
-		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.m", "org.z", true))
+		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, utils.CreateStringPartition("com.m", "org.z", true).Range)
 		require.NoError(t, err)
 		require.Equal(t, bson.D{
 			bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{

flow/connectors/mysql/qrep_partition.go

@@ -8,14 +8,11 @@ import (
 
 	"github.com/google/uuid"
 
+	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
 	"github.com/PeerDB-io/peerdb/flow/shared"
 )
 
-// String watermark partitioning is currently only supported by the MySQL connector,
-// so these helpers live in the MySQL connector subtree to prevent accidental usage
-// by other connectors.
-
 var (
 	uuidLowerRe = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`)
 	uuidUpperRe = regexp.MustCompile(`^[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}$`)
@@ -76,10 +73,10 @@ func buildUuidStringPartitions(
 		if err != nil {
 			return nil, fmt.Errorf("failed to convert bigint to uuid: %w", err)
 		}
-		partitions = append(partitions, createStringPartition(start, end, false))
+		partitions = append(partitions, utils.CreateStringPartition(start, end, false))
 		start = end
 	}
-	partitions = append(partitions, createStringPartition(start, maxVal, true))
+	partitions = append(partitions, utils.CreateStringPartition(start, maxVal, true))
 	return partitions, nil
 }
 
@@ -103,18 +100,3 @@ 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,
-				},
-			},
-		},
-	}
-}

flow/connectors/utils/monitoring/monitoring.go

@@ -340,6 +340,8 @@ func addPartitionToQRepRun(ctx context.Context, tx pgx.Tx, flowJobName string,
 		switch x := partition.Range.Range.(type) {
 		case *protos.PartitionRange_IntRange:
 			rangeStart, rangeEnd = new(strconv.FormatInt(x.IntRange.Start, 10)), new(strconv.FormatInt(x.IntRange.End, 10))
+		case *protos.PartitionRange_NumericRange:
+			rangeStart, rangeEnd = new(strconv.FormatInt(x.NumericRange.Start, 10)), new(strconv.FormatInt(x.NumericRange.End, 10))
 		case *protos.PartitionRange_UintRange:
 			rangeStart, rangeEnd = new(strconv.FormatUint(x.UintRange.Start, 10)), new(strconv.FormatUint(x.UintRange.End, 10))
 		case *protos.PartitionRange_TimestampRange:

flow/connectors/utils/partition.go

@@ -409,3 +409,33 @@ func (p *PartitionHelper) AddPartitions(partitions []*protos.QRepPartition) {
 func (p *PartitionHelper) GetPartitions() []*protos.QRepPartition {
 	return p.partitions
 }
+
+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,
+				},
+			},
+		},
+	}
+}
+
+func CreateNumericPartition(start int64, end int64, endInclusive bool) *protos.QRepPartition {
+	return &protos.QRepPartition{
+		PartitionId: uuid.NewString(),
+		Range: &protos.PartitionRange{
+			Range: &protos.PartitionRange_NumericRange{
+				NumericRange: &protos.NumericPartitionRange{
+					Start:        start,
+					End:          end,
+					EndInclusive: endInclusive,
+				},
+			},
+		},
+	}
+}