using the full 12 bytes instead of just first 4 bytes of objectID for partitioning

Author: jgao54

flow/connectors/mongo/qrep_partition.go

@@ -2,11 +2,9 @@ package connmongo
 
 import (
 	"context"
-	"encoding/binary"
 	"fmt"
 	"log/slog"
-	"math"
-	"time"
+	"math/big"
 
 	"github.com/google/uuid"
 	"go.mongodb.org/mongo-driver/v2/bson"
@@ -18,29 +16,12 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/shared"
 )
 
-// minObjectIDForTimestamp creates an ObjectID with the given timestamp and
-// all trailing bytes set to 0x00, producing the smallest possible ObjectID
-// for that second.
-func minObjectIDForTimestamp(t time.Time) bson.ObjectID {
-	var oid [12]byte
-	binary.BigEndian.PutUint32(oid[0:4], uint32(t.Unix()))
-	return oid
-}
-
-// maxObjectIDForTimestamp creates an ObjectID with the given timestamp and
-// all trailing bytes set to 0xFF (max unsigned 64-bit value), producing
-// the largest possible ObjectID for that second.
-func maxObjectIDForTimestamp(t time.Time) bson.ObjectID {
-	var oid [12]byte
-	binary.BigEndian.PutUint32(oid[0:4], uint32(t.Unix()))
-	binary.BigEndian.PutUint64(oid[4:], math.MaxUint64)
-	return oid
-}
-
 // minMaxPartitions creates partitions by querying only the min and max _id
 // from the collection (leveraging the default _id index), then uniformly
-// dividing the timestamp range encoded in the ObjectIDs. Note that partition
-// may be skewed if document insertion rate varies significantly over time.
+// 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,
@@ -61,34 +42,44 @@ func (c *MongoConnector) minMaxPartitions(
 
 	minID := minRaw.ObjectID()
 	maxID := maxRaw.ObjectID()
-	minTs := minID.Timestamp()
-	maxTs := maxID.Timestamp()
-	tsRange := maxTs.Unix() - minTs.Unix() + 1
-	if tsRange <= 1 {
-		c.logger.Info("[mongo] min/max timestamps are equal, falling back to full table partition")
+
+	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 ObjectIDs are equal, falling back to full table partition")
 		return fullTablePartition, nil
 	}
-	// Cap partitions to the timestamp range in seconds so we don't create additional
-	// empty partitions unnecessarily when docs span fewer seconds than partitions
-	numPartitions = min(numPartitions, tsRange)
 
-	c.logger.Info("[mongo] using min/max ObjectID timestamp partitioning",
-		slog.Time("minTimestamp", minTs),
-		slog.Time("maxTimestamp", maxTs),
+	c.logger.Info("[mongo] using min/max ObjectID partitioning",
+		slog.String("minID", minID.Hex()),
+		slog.String("maxID", maxID.Hex()),
 		slog.Int64("numPartitions", numPartitions))
 
-	secondsPerPartition := shared.DivCeil(tsRange, numPartitions)
+	step := shared.BigIntDivCeil(intRange, big.NewInt(numPartitions))
+
 	partitions := make([]*protos.QRepPartition, 0, numPartitions)
-	for i := range numPartitions {
-		start := minObjectIDForTimestamp(time.Unix(minTs.Unix()+secondsPerPartition*i, 0))
-		end := maxObjectIDForTimestamp(time.Unix(minTs.Unix()+secondsPerPartition*(i+1)-1, 0))
+	for start := new(big.Int).Set(minInt); start.Cmp(maxInt) <= 0; start.Add(start, step) {
+		end := new(big.Int).Add(start, step)
+		end.Sub(end, big.NewInt(1))
+		if end.Cmp(maxInt) > 0 {
+			end.Set(maxInt)
+		}
+		startObjectID, err := bigIntToObjectID(start)
+		if err != nil {
+			return nil, err
+		}
+		endObjectID, err := bigIntToObjectID(end)
+		if err != nil {
+			return nil, err
+		}
 		partitions = append(partitions, &protos.QRepPartition{
 			PartitionId: uuid.NewString(),
 			Range: &protos.PartitionRange{
 				Range: &protos.PartitionRange_ObjectIdRange{
 					ObjectIdRange: &protos.ObjectIdPartitionRange{
-						Start: start.Hex(),
-						End:   end.Hex(),
+						Start: startObjectID.Hex(),
+						End:   endObjectID.Hex(),
 					},
 				},
 			},
@@ -99,6 +90,16 @@ 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)")
+	}
+	n.FillBytes(oid[:])
+	return oid, 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

flow/shared/math.go

@@ -1,13 +1,21 @@
 package shared
 
 import (
+	"math/big"
+
 	"golang.org/x/exp/constraints"
 )
 
 func DivCeil[T constraints.Integer](x, y T) T {
 	return (x + y - 1) / y
 }
 
+// BigIntDivCeil returns ceil(x / y) for big.Int values.
+func BigIntDivCeil(x, y *big.Int) *big.Int {
+	result := new(big.Int).Add(x, new(big.Int).Sub(y, big.NewInt(1)))
+	return result.Div(result, y)
+}
+
 // AdjustedPartitions represents the adjusted partitioning parameters
 type AdjustedPartitions struct {
 	AdjustedNumPartitions       int64