mongo: address PR review comments on _id partitioning

- Add end_inclusive to StringPartitionRange proto (aligns with MySQL PR
  #4482 adding the same field; avoids future merge conflict)
- Replace sentinel-string approach with end_inclusive=true on the last
  string partition, using $lte instead of $lt for its MongoDB filter
- Sort string _id samples via MongoDB $sort stage instead of Go sort,
  preserving the collection's collation order
- Filter boundary samples by exact equality instead of <=/>= comparison,
  trusting MongoDB's min/max guarantees
- Add e2e tests for string and numeric _id partitioned snapshot flows
- Document the double _id gap limitation on numericPartitions (follow-up)

Author: andreyzhelnin-st

flow/connectors/mongo/qrep.go

@@ -245,13 +245,16 @@ func toRangeFilter(watermarkColumn string, partitionRange *protos.PartitionRange
 			}},
 		}, 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.
+		// 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"
+		}
 		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},
+				bson.E{Key: endOp, Value: r.StringRange.End},
 			}},
 		}, nil
 	default:

flow/connectors/mongo/qrep_partition.go

@@ -17,14 +17,6 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/shared"
 )
 
-// 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.
@@ -150,6 +142,12 @@ func (c *MongoConnector) objectIDPartitions(
 // 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.
 func (c *MongoConnector) numericPartitions(
 	minVal int64,
 	maxVal int64,
@@ -207,14 +205,15 @@ func (c *MongoConnector) stringPartitions(
 		slog.Int("sampleCount", len(samples)))
 
 	partitions := make([]*protos.QRepPartition, 0, len(ranges))
-	for _, rng := range 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],
+						Start:        rng[0],
+						End:          rng[1],
+						EndInclusive: i == len(ranges)-1,
 					},
 				},
 			},
@@ -242,6 +241,7 @@ 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}}}},
 		}},
 		{Key: "cursor", Value: bson.D{}},
@@ -267,20 +267,26 @@ func (c *MongoConnector) sampleStringIDs(
 	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.
+// computeStringBoundaries turns a pre-sorted (by database collation) sample of
+// string _ids plus the real min/max into a contiguous set of ranges. Interior
+// boundaries are quantiles of the (deduplicated) sample; the first range starts
+// at the real min and the last range ends at the real max. All ranges except the
+// last are half-open [start, end); the last is closed [start, end].
+//
+// samples must already be sorted in the database's collation order; this function
+// does not re-sort them, to preserve collation semantics.
 //
 // 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).
+	// Keep unique sampled values strictly inside (min, max), preserving order.
+	// We use exact equality to filter boundaries: MongoDB guarantees all sampled
+	// _ids are within [minVal, maxVal] by its own collation; only exact matches
+	// with the boundary values would create zero-width partitions.
 	seen := make(map[string]struct{}, len(samples))
 	interior := make([]string, 0, len(samples))
 	for _, s := range samples {
-		if s <= minVal || s >= maxVal {
+		if s == minVal || s == maxVal {
 			continue
 		}
 		if _, ok := seen[s]; ok {
@@ -289,7 +295,6 @@ func computeStringBoundaries(minVal string, maxVal string, samples []string, num
 		seen[s] = struct{}{}
 		interior = append(interior, s)
 	}
-	slices.Sort(interior)
 
 	// Pick up to numPartitions-1 evenly spaced interior boundaries (quantiles).
 	desiredBoundaries := numPartitions - 1
@@ -308,15 +313,16 @@ func computeStringBoundaries(minVal string, maxVal string, samples []string, num
 		picked = slices.Compact(picked) // drop consecutive duplicate quantiles
 	}
 
-	// Build contiguous half-open ranges: [min, p0), [p0, p1), ..., [pLast, max+sentinel).
+	// Build ranges: [min, p0), [p0, p1), ..., [pLast-1, pLast), [pLast, max].
+	// The last range is closed (end_inclusive=true); all others are half-open.
 	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
+			end = maxVal
 		}
 		ranges[i] = [2]string{starts[i], end}
 	}

flow/connectors/mongo/qrep_partition_test.go

@@ -44,13 +44,13 @@ func TestIsNumericIDType(t *testing.T) {
 }
 
 // assertContiguous verifies the ranges form a gap-free, non-overlapping cover of
-// [min, max+sentinel): first start == min, last end == max+sentinel, and each
-// range's end equals the next range's start.
+// [min, max]: first start == min, last end == max, and each range's end equals
+// the next range's start (half-open except for the last which is closed).
 func assertContiguous(t *testing.T, ranges [][2]string, minVal, maxVal string) {
 	t.Helper()
 	require.NotEmpty(t, ranges)
 	require.Equal(t, minVal, ranges[0][0], "first partition must start at min")
-	require.Equal(t, maxVal+stringPartitionUpperSentinel, ranges[len(ranges)-1][1], "last partition must end just past max")
+	require.Equal(t, maxVal, ranges[len(ranges)-1][1], "last partition must end at max")
 	for i := range ranges {
 		require.Less(t, ranges[i][0], ranges[i][1], "range start must be < end")
 		if i+1 < len(ranges) {
@@ -88,9 +88,10 @@ func TestComputeStringBoundaries(t *testing.T) {
 		assertContiguous(t, ranges, "a", "z")
 	})
 
-	t.Run("samples outside (min,max) ignored", func(t *testing.T) {
-		// "a" == min and "z" == max must be dropped; only "m" is a valid interior boundary
-		samples := []string{"a", "z", "m", "0", "zzzz"}
+	t.Run("boundary samples ignored", func(t *testing.T) {
+		// "a" == min and "z" == max must be dropped to avoid zero-width partitions;
+		// only "m" is a valid interior boundary. Samples come pre-sorted from MongoDB.
+		samples := []string{"a", "m", "z"}
 		ranges := computeStringBoundaries("a", "z", samples, 4)
 		require.Len(t, ranges, 2)
 		assertContiguous(t, ranges, "a", "z")
@@ -122,9 +123,9 @@ func intPartition(start, end int64) *protos.PartitionRange {
 	}}
 }
 
-func stringPartition(start, end string) *protos.PartitionRange {
+func stringPartition(start, end string, endInclusive bool) *protos.PartitionRange {
 	return &protos.PartitionRange{Range: &protos.PartitionRange_StringRange{
-		StringRange: &protos.StringPartitionRange{Start: start, End: end},
+		StringRange: &protos.StringPartitionRange{Start: start, End: end, EndInclusive: endInclusive},
 	}}
 }
 
@@ -141,7 +142,7 @@ func TestToRangeFilter(t *testing.T) {
 	})
 
 	t.Run("string range is half-open", func(t *testing.T) {
-		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.a", "com.m"))
+		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.a", "com.m", false))
 		require.NoError(t, err)
 		require.Equal(t, bson.D{
 			bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{
@@ -151,6 +152,17 @@ func TestToRangeFilter(t *testing.T) {
 		}, filter)
 	})
 
+	t.Run("last string range is closed", func(t *testing.T) {
+		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.m", "org.z", true))
+		require.NoError(t, err)
+		require.Equal(t, bson.D{
+			bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{
+				bson.E{Key: "$gte", Value: "com.m"},
+				bson.E{Key: "$lte", Value: "org.z"},
+			}},
+		}, filter)
+	})
+
 	t.Run("object id range still supported", func(t *testing.T) {
 		r := &protos.PartitionRange{Range: &protos.PartitionRange_ObjectIdRange{
 			ObjectIdRange: &protos.ObjectIdPartitionRange{

flow/e2e/mongo_test.go

@@ -157,6 +157,120 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned() {
 	RequireEnvCanceled(t, env)
 }
 
+func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_StringID() {
+	t := s.T()
+	srcDatabase := GetTestDatabase(s.Suffix())
+	srcTable := "test_simple_partitioned_string_id"
+	dstTable := "test_simple_dst_partitioned_string_id"
+
+	connectionGen := FlowConnectionGenerationConfig{
+		FlowJobName:   AddSuffix(s, srcTable),
+		TableMappings: TableMappings(s, srcTable, dstTable),
+		Destination:   s.Peer().Name,
+	}
+	flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen)
+	flowConnConfig.DoInitialSnapshot = true
+	flowConnConfig.SnapshotNumRowsPerPartition = 10
+
+	adminClient := s.Source().(*MongoSource).AdminClient()
+	collection := adminClient.Database(srcDatabase).Collection(srcTable)
+
+	docs := make([]any, 1000)
+	for i := range 1000 {
+		docs[i] = bson.D{
+			{Key: "_id", Value: fmt.Sprintf("id-%05d", i)},
+			{Key: fmt.Sprintf("init_key_%d", i), Value: fmt.Sprintf("init_value_%d", i)},
+		}
+	}
+	_, err := collection.InsertMany(t.Context(), docs)
+	require.NoError(t, err)
+
+	tc := NewTemporalClient(t)
+	env := ExecutePeerflow(t, tc, flowConnConfig)
+
+	EnvWaitForEqualTablesWithNames(env, s, "initial load to match", srcTable, dstTable, "_id,doc")
+
+	catalogPool, err := internal.GetCatalogConnectionPoolFromEnv(t.Context())
+	require.NoError(t, err)
+	var partitionCount int
+	require.NoError(t, catalogPool.QueryRow(t.Context(),
+		`SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`,
+		flowConnConfig.FlowJobName).Scan(&partitionCount))
+	require.Greater(t, partitionCount, 1, "expected multiple partitions for 1000 string-id rows")
+
+	SetupCDCFlowStatusQuery(t, env, flowConnConfig)
+	cdcDocs := make([]any, 10)
+	for i := range 10 {
+		cdcDocs[i] = bson.D{
+			{Key: "_id", Value: fmt.Sprintf("cdc-%05d", i)},
+			{Key: fmt.Sprintf("cdc_key_%d", i), Value: fmt.Sprintf("cdc_value_%d", i)},
+		}
+	}
+	_, err = collection.InsertMany(t.Context(), cdcDocs)
+	require.NoError(t, err)
+
+	EnvWaitForEqualTablesWithNames(env, s, "cdc events to match", srcTable, dstTable, "_id,doc")
+	env.Cancel(t.Context())
+	RequireEnvCanceled(t, env)
+}
+
+func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_NumericID() {
+	t := s.T()
+	srcDatabase := GetTestDatabase(s.Suffix())
+	srcTable := "test_simple_partitioned_numeric_id"
+	dstTable := "test_simple_dst_partitioned_numeric_id"
+
+	connectionGen := FlowConnectionGenerationConfig{
+		FlowJobName:   AddSuffix(s, srcTable),
+		TableMappings: TableMappings(s, srcTable, dstTable),
+		Destination:   s.Peer().Name,
+	}
+	flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen)
+	flowConnConfig.DoInitialSnapshot = true
+	flowConnConfig.SnapshotNumRowsPerPartition = 10
+
+	adminClient := s.Source().(*MongoSource).AdminClient()
+	collection := adminClient.Database(srcDatabase).Collection(srcTable)
+
+	docs := make([]any, 1000)
+	for i := range 1000 {
+		docs[i] = bson.D{
+			{Key: "_id", Value: int32(i + 1)},
+			{Key: fmt.Sprintf("init_key_%d", i), Value: fmt.Sprintf("init_value_%d", i)},
+		}
+	}
+	_, err := collection.InsertMany(t.Context(), docs)
+	require.NoError(t, err)
+
+	tc := NewTemporalClient(t)
+	env := ExecutePeerflow(t, tc, flowConnConfig)
+
+	EnvWaitForEqualTablesWithNames(env, s, "initial load to match", srcTable, dstTable, "_id,doc")
+
+	catalogPool, err := internal.GetCatalogConnectionPoolFromEnv(t.Context())
+	require.NoError(t, err)
+	var partitionCount int
+	require.NoError(t, catalogPool.QueryRow(t.Context(),
+		`SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`,
+		flowConnConfig.FlowJobName).Scan(&partitionCount))
+	require.Greater(t, partitionCount, 1, "expected multiple partitions for 1000 numeric-id rows")
+
+	SetupCDCFlowStatusQuery(t, env, flowConnConfig)
+	cdcDocs := make([]any, 10)
+	for i := range 10 {
+		cdcDocs[i] = bson.D{
+			{Key: "_id", Value: int32(10000 + i)},
+			{Key: fmt.Sprintf("cdc_key_%d", i), Value: fmt.Sprintf("cdc_value_%d", i)},
+		}
+	}
+	_, err = collection.InsertMany(t.Context(), cdcDocs)
+	require.NoError(t, err)
+
+	EnvWaitForEqualTablesWithNames(env, s, "cdc events to match", srcTable, dstTable, "_id,doc")
+	env.Cancel(t.Context())
+	RequireEnvCanceled(t, env)
+}
+
 func (s MongoClickhouseSuite) Test_Snapshot_Collection_With_Single_Document() {
 	t := s.T()
 	srcDatabase := GetTestDatabase(s.Suffix())

protos/flow.proto

@@ -327,6 +327,10 @@ message ObjectIdPartitionRange {
 message StringPartitionRange {
   string start = 1;
   string end = 2;
+  // unlike numeric/temporal ranges, string boundaries cannot be safely
+  // incremented; so the partition boundaries are [start, end) except
+  // for the last partition being [start, end]; hence this bool is needed.
+  bool end_inclusive = 3;
 }
 
 message PartitionRange {