mongo: partition snapshots for numeric and string _id collections

The MongoDB connector only parallelized initial snapshots when _id was a
native ObjectID; numeric or string _id collections fell back to a single
serial full-table partition. For large collections keyed by natural ids
(e.g. app-store numeric ids or package-name strings) this makes the
snapshot too slow to finish within the oplog retention window.

Dispatch in GetQRepPartitions on the collection's min/max _id BSON type:
- ObjectID: unchanged (uniform division of the 12-byte keyspace).
- Int32/Int64: uniform min/max division, reusing utils.PartitionHelper
  to emit non-overlapping IntPartitionRange partitions.
- String: sampling-based quantile boundaries ($sample, honouring the
  configured read preference) anchored by the real min/max, emitted as a
  new StringPartitionRange proto type with contiguous half-open ranges.
Mixed-type, Double, or empty collections fall back to full-table.

Add IntRange/StringRange cases to toRangeFilter and StringRange handling
to monitoring.addPartitionToQRepRun. computeStringBoundaries is pure and
unit-tested alongside the range-to-filter conversions.

Author: andreyzhelnin-st

flow/connectors/mongo/qrep.go

@@ -73,7 +73,7 @@ func (c *MongoConnector) GetQRepPartitions(
 		return fullTablePartition, nil
 	}
 
-	return c.minMaxPartitions(ctx, collection, adjustedPartitions.AdjustedNumPartitions)
+	return c.buildPartitions(ctx, collection, adjustedPartitions.AdjustedNumPartitions)
 }
 
 func (c *MongoConnector) GetDefaultPartitionKeyForTables(
@@ -234,6 +234,26 @@ 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.
+		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},
+			}},
+		}, nil
+	case *protos.PartitionRange_StringRange:
+		// String _id: half-open [start, end) range. computeStringBoundaries makes
+		// ranges contiguous and sets the final end just past the real max, so $lt
+		// gives exact coverage with no gaps or overlap.
+		return bson.D{
+			bson.E{Key: watermarkColumn, Value: bson.D{
+				bson.E{Key: "$gte", Value: r.StringRange.Start},
+				bson.E{Key: "$lt", Value: r.StringRange.End},
+			}},
+		}, nil
 	default:
 		return nil, fmt.Errorf("unsupported partition range type")
 	}

flow/connectors/mongo/qrep_partition.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log/slog"
 	"math/big"
+	"slices"
 
 	"github.com/google/uuid"
 	"go.mongodb.org/mongo-driver/v2/bson"
@@ -16,37 +17,95 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/shared"
 )
 
-// minMaxPartitions creates partitions by querying only the min and max _id
-// from the collection (leveraging the default _id index), then uniformly
-// dividing the ObjectID range using integer arithmetic on the full 12-byte
-// value. Since the first 4-byte timestamp is the most significant component,
-// this naturally partitions by time while also provides best-effort handling
-// of edge cases where records are all inserted in the same second.
-func (c *MongoConnector) minMaxPartitions(
-	ctx context.Context,
-	collection *mongo.Collection,
-	numPartitions int64,
-) ([]*protos.QRepPartition, error) {
-	fullTablePartition := []*protos.QRepPartition{{
+// stringPartitionUpperSentinel is appended to the collection's max string _id to
+// form the exclusive upper bound of the final string partition. It is the highest
+// valid Unicode code point, so `max + sentinel` sorts strictly after `max` (and
+// after every real _id, since `max` is the maximum). This lets every string
+// partition use uniform half-open [start, end) ($gte/$lt) semantics while still
+// capturing the maximum document in the last partition.
+const stringPartitionUpperSentinel = "\U0010FFFF"
+
+const (
+	// stringSampleOversample draws more samples than partitions so quantile
+	// boundaries are well distributed even with clustered keys.
+	stringSampleOversample = 20
+	// stringSampleMaxSize caps sampling cost on very large collections.
+	stringSampleMaxSize = 100000
+)
+
+// fullTablePartitions returns the single-partition full-table fallback used when a
+// collection cannot be partitioned (unsupported/mixed _id type, empty collection,
+// or degenerate min/max range).
+func fullTablePartitions() []*protos.QRepPartition {
+	return []*protos.QRepPartition{{
 		PartitionId:        utils.FullTablePartitionID,
 		Range:              nil,
 		FullTablePartition: true,
 	}}
+}
 
-	minID, maxID, found, err := findMinMaxObjectIDs(ctx, collection, c.config.ReadPreference)
+// 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)
+//   - 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.
+func (c *MongoConnector) buildPartitions(
+	ctx context.Context,
+	collection *mongo.Collection,
+	numPartitions int64,
+) ([]*protos.QRepPartition, error) {
+	minRaw, err := findBoundaryID(ctx, collection, Lower, c.config.ReadPreference)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to find min _id: %w", err)
+	}
+	if minRaw.IsZero() {
+		c.logger.Info("[mongo] no documents found, falling back to full table partition")
+		return fullTablePartitions(), nil
 	}
-	if !found {
-		c.logger.Info("[mongo] no valid min/max ObjectIDs found, falling back to full table partition")
-		return fullTablePartition, nil
+	maxRaw, err := findBoundaryID(ctx, collection, Upper, c.config.ReadPreference)
+	if err != nil {
+		return nil, fmt.Errorf("failed to find max _id: %w", err)
+	}
+	if maxRaw.IsZero() {
+		c.logger.Info("[mongo] no documents found, falling back to full table partition")
+		return fullTablePartitions(), nil
 	}
+
+	switch {
+	case minRaw.Type == bson.TypeObjectID && maxRaw.Type == bson.TypeObjectID:
+		return c.objectIDPartitions(minRaw.ObjectID(), maxRaw.ObjectID(), numPartitions)
+	case isNumericIDType(minRaw.Type) && isNumericIDType(maxRaw.Type):
+		minVal, _ := rawValueToInt64(minRaw)
+		maxVal, _ := rawValueToInt64(maxRaw)
+		return c.numericPartitions(minVal, maxVal, numPartitions)
+	case minRaw.Type == bson.TypeString && maxRaw.Type == bson.TypeString:
+		return c.stringPartitions(ctx, collection, minRaw.StringValue(), maxRaw.StringValue(), numPartitions)
+	default:
+		c.logger.Info("[mongo] _id type not supported for partitioning (or mixed types), falling back to full table partition",
+			slog.String("minType", minRaw.Type.String()),
+			slog.String("maxType", maxRaw.Type.String()))
+		return fullTablePartitions(), nil
+	}
+}
+
+// objectIDPartitions divides the ObjectID keyspace uniformly using integer
+// arithmetic on the full 12-byte value. Since the leading 4-byte timestamp is the
+// most significant component, this naturally partitions by insertion time.
+func (c *MongoConnector) objectIDPartitions(
+	minID bson.ObjectID,
+	maxID bson.ObjectID,
+	numPartitions int64,
+) ([]*protos.QRepPartition, error) {
 	minInt := new(big.Int).SetBytes(minID[:])
 	maxInt := new(big.Int).SetBytes(maxID[:])
 	intRange := new(big.Int).Sub(maxInt, minInt)
 	if intRange.Sign() <= 0 {
 		c.logger.Info("[mongo] min/max ObjectID range is non-positive, falling back to full table partition")
-		return fullTablePartition, nil
+		return fullTablePartitions(), nil
 	}
 
 	c.logger.Info("[mongo] using min/max ObjectID partitioning",
@@ -88,43 +147,205 @@ func (c *MongoConnector) minMaxPartitions(
 	return partitions, nil
 }
 
-// bigIntToObjectID converts a big.Int back to a 12-byte ObjectID.
-func bigIntToObjectID(n *big.Int) (bson.ObjectID, error) {
-	var oid bson.ObjectID
-	if n.BitLen() > 96 {
-		return oid, fmt.Errorf("big.Int value exceeds 96 bits (ObjectID size)")
+// 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.
+func (c *MongoConnector) numericPartitions(
+	minVal int64,
+	maxVal int64,
+	numPartitions int64,
+) ([]*protos.QRepPartition, error) {
+	if maxVal <= minVal {
+		c.logger.Info("[mongo] numeric min/max range is non-positive, falling back to full table partition")
+		return fullTablePartitions(), nil
 	}
-	n.FillBytes(oid[:])
-	return oid, nil
+
+	c.logger.Info("[mongo] using numeric _id partitioning",
+		slog.Int64("minID", minVal),
+		slog.Int64("maxID", maxVal),
+		slog.Int64("numPartitions", numPartitions))
+
+	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)
+	}
+	return helper.GetPartitions(), nil
 }
 
-// findMinMaxObjectIDs returns the min and max object IDs from the collection.
-// The two boundary queries are not wrapped in a snapshot session: a concurrent
-// insert between them may shift the true min/max, but this is safe because the
-// change stream resume token is captured before either query runs. Any writes
-// that occur during or after partitioning will be replayed by the change stream.
-func findMinMaxObjectIDs(
+// stringPartitions builds partitions for a string _id collection. Because string
+// keys (e.g. package names) are not uniformly distributed, uniform division of the
+// keyspace would be badly skewed; instead we sample the collection and pick quantile
+// boundaries so each partition holds a roughly equal share of documents.
+func (c *MongoConnector) stringPartitions(
 	ctx context.Context,
 	collection *mongo.Collection,
-	readPreference protos.ReadPreference,
-) (bson.ObjectID, bson.ObjectID, bool, error) {
-	minRaw, err := findBoundaryID(ctx, collection, Lower, readPreference)
+	minVal string,
+	maxVal string,
+	numPartitions int64,
+) ([]*protos.QRepPartition, error) {
+	if minVal >= maxVal {
+		c.logger.Info("[mongo] string min/max range is non-positive, falling back to full table partition")
+		return fullTablePartitions(), nil
+	}
+
+	samples, err := c.sampleStringIDs(ctx, collection, numPartitions)
 	if err != nil {
-		return bson.ObjectID{}, bson.ObjectID{}, false, fmt.Errorf("failed to find min _id: %w", err)
+		return nil, err
+	}
+
+	ranges := computeStringBoundaries(minVal, maxVal, samples, numPartitions)
+	if len(ranges) < 2 {
+		c.logger.Info("[mongo] insufficient string boundaries from sampling, falling back to full table partition",
+			slog.Int("sampleCount", len(samples)))
+		return fullTablePartitions(), nil
+	}
+
+	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 _, 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],
+					},
+				},
+			},
+			FullTablePartition: false,
+		})
 	}
-	if minRaw.Type != bson.TypeObjectID {
-		return bson.ObjectID{}, bson.ObjectID{}, false, nil
+	return partitions, nil
+}
+
+// sampleStringIDs draws a random sample of string _id values. $sample with a size
+// below ~5% of the collection uses WiredTiger's random cursor, so this is cheap even
+// on large collections; it honours the configured read preference (run on a
+// secondary to avoid loading the primary).
+func (c *MongoConnector) sampleStringIDs(
+	ctx context.Context,
+	collection *mongo.Collection,
+	numPartitions int64,
+) ([]string, error) {
+	sampleSize := numPartitions * stringSampleOversample
+	if sampleSize > stringSampleMaxSize {
+		sampleSize = stringSampleMaxSize
 	}
 
-	maxRaw, err := findBoundaryID(ctx, collection, Upper, readPreference)
+	aggCmd := bson.D{
+		{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: "$project", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}}},
+		}},
+		{Key: "cursor", Value: bson.D{}},
+	}
+
+	cursor, err := collection.Database().RunCommandCursor(ctx, aggCmd,
+		options.RunCmd().SetReadPreference(protoToReadPref[c.config.ReadPreference]))
 	if err != nil {
-		return bson.ObjectID{}, bson.ObjectID{}, false, fmt.Errorf("failed to find max _id: %w", err)
+		return nil, fmt.Errorf("failed to sample _id values: %w", err)
+	}
+	defer cursor.Close(ctx)
+
+	samples := make([]string, 0, sampleSize)
+	for cursor.Next(ctx) {
+		rv := cursor.Current.Lookup(DefaultDocumentKeyColumnName)
+		if rv.Type == bson.TypeString {
+			samples = append(samples, rv.StringValue())
+		}
+	}
+	if err := cursor.Err(); err != nil {
+		return nil, fmt.Errorf("cursor error while sampling _id values: %w", err)
+	}
+	return samples, nil
+}
+
+// computeStringBoundaries turns a random sample of string _ids plus the real
+// min/max into a contiguous set of half-open [start, end) ranges. Interior
+// boundaries are quantiles of the (deduplicated, sorted) sample; the first range
+// starts at the real min and the last range ends just past the real max (via the
+// sentinel) so coverage of [min, max] is exact with no gaps or overlap.
+//
+// It is pure (no I/O) to keep the boundary math unit-testable. Returns fewer ranges
+// than numPartitions when the sample yields too few distinct interior boundaries.
+func computeStringBoundaries(minVal string, maxVal string, samples []string, numPartitions int64) [][2]string {
+	// Keep unique sampled values strictly inside (min, max).
+	seen := make(map[string]struct{}, len(samples))
+	interior := make([]string, 0, len(samples))
+	for _, s := range samples {
+		if s <= minVal || s >= maxVal {
+			continue
+		}
+		if _, ok := seen[s]; ok {
+			continue
+		}
+		seen[s] = struct{}{}
+		interior = append(interior, s)
+	}
+	slices.Sort(interior)
+
+	// Pick up to numPartitions-1 evenly spaced interior boundaries (quantiles).
+	desiredBoundaries := numPartitions - 1
+	var picked []string
+	if int64(len(interior)) <= desiredBoundaries {
+		picked = interior
+	} else {
+		picked = make([]string, 0, desiredBoundaries)
+		for i := int64(1); i <= desiredBoundaries; i++ {
+			idx := i * int64(len(interior)) / numPartitions
+			if idx >= int64(len(interior)) {
+				idx = int64(len(interior)) - 1
+			}
+			picked = append(picked, interior[idx])
+		}
+		picked = slices.Compact(picked) // drop consecutive duplicate quantiles
 	}
-	if maxRaw.Type != bson.TypeObjectID {
-		return bson.ObjectID{}, bson.ObjectID{}, false, nil
+
+	// Build contiguous half-open ranges: [min, p0), [p0, p1), ..., [pLast, max+sentinel).
+	starts := append([]string{minVal}, picked...)
+	ranges := make([][2]string, len(starts))
+	for i := range starts {
+		var end string
+		if i+1 < len(starts) {
+			end = starts[i+1]
+		} else {
+			end = maxVal + stringPartitionUpperSentinel
+		}
+		ranges[i] = [2]string{starts[i], end}
 	}
+	return ranges
+}
 
-	return minRaw.ObjectID(), maxRaw.ObjectID(), true, nil
+// bigIntToObjectID converts a big.Int back to a 12-byte ObjectID.
+func bigIntToObjectID(n *big.Int) (bson.ObjectID, error) {
+	var oid bson.ObjectID
+	if n.BitLen() > 96 {
+		return oid, fmt.Errorf("big.Int value exceeds 96 bits (ObjectID size)")
+	}
+	n.FillBytes(oid[:])
+	return oid, nil
+}
+
+func isNumericIDType(t bson.Type) bool {
+	return t == bson.TypeInt32 || t == bson.TypeInt64
+}
+
+func rawValueToInt64(rv bson.RawValue) (int64, bool) {
+	switch rv.Type {
+	case bson.TypeInt32:
+		return int64(rv.Int32()), true
+	case bson.TypeInt64:
+		return rv.Int64(), true
+	default:
+		return 0, false
+	}
 }
 
 type Boundary int