Update getMultiTypeIDKeyRanges function to exclute NaN _id values.

Author: inelpandzic

plm/copy.go

@@ -654,24 +654,25 @@ func NewSegmenter(
 		return s, nil
 	}
 
-	keyRangeByType, err := getIDKeyRangeByType(ctx, mcoll)
+	multiTypeIDkeyRanges, err := getMultiTypeIDKeyRanges(ctx, mcoll)
 	if err != nil {
 		return nil, errors.Wrap(err, "get ID key range by type")
 	}
 
-	if len(keyRangeByType) == 0 {
+	if len(multiTypeIDkeyRanges) == 0 {
 		return nil, errEOC // empty collection
 	}
 
-	currIDRange := keyRangeByType[0]
-	keyRanges := keyRangeByType[1:]
+	currIDRange := multiTypeIDkeyRanges[0]
+	remainingKeyRanges := multiTypeIDkeyRanges[1:]
 
 	s := &Segmenter{
 		mcoll:       mcoll,
 		segmentSize: segmentSize,
 		batchSize:   batchSize,
-		keyRanges:   keyRanges,
+		keyRanges:   remainingKeyRanges,
 		currIDRange: currIDRange,
+		nanDoc:      *nanDoc,
 	}
 
 	return s, nil
@@ -848,30 +849,44 @@ func getIDKeyRange(ctx context.Context, mcoll *mongo.Collection) (keyRange, *bso
 	return ret, &nanDoc, nil
 }
 
-// getIDKeyRangeByType returns a slice of keyRange grouped by the BSON type of the _id field.
+// getMultiTypeIDKeyRanges returns a slice of keyRange grouped by the BSON type of the _id field.
 // It performs an aggregation that groups documents by _id type, computing the min and max _id
 // for each group. This allows the Segmenter to handle collections with heterogeneous _id types
 // by processing each type range sequentially.
-func getIDKeyRangeByType(ctx context.Context, mcoll *mongo.Collection) ([]keyRange, error) {
-	cur, err := mcoll.Aggregate(ctx, mongo.Pipeline{
-		bson.D{{"$group", bson.D{
-			{"_id", bson.D{{"type", bson.D{{"$type", "$_id"}}}}},
-			{"minKey", bson.D{{"$min", "$_id"}}},
-			{"maxKey", bson.D{{"$max", "$_id"}}},
-		}}},
-	})
+func getMultiTypeIDKeyRanges(ctx context.Context, mcoll *mongo.Collection) ([]keyRange, error) {
+	cur, err := mcoll.Aggregate(ctx,
+		mongo.Pipeline{
+			// Match only numeric types that are not NaN
+			bson.D{{"$match", bson.D{
+				{"$expr", bson.D{
+					// Only allow if _id is not NaN
+					{"$ne", bson.A{"$_id", bson.D{{"$literal", math.NaN()}}}},
+				}},
+			}}},
+			// Group by type and find min/max
+			bson.D{{"$group", bson.D{
+				{"_id", bson.D{{"type", bson.D{{"$type", "$_id"}}}}},
+				{"minKey", bson.D{{"$min", "$_id"}}},
+				{"maxKey", bson.D{{"$max", "$_id"}}},
+			}}},
+		})
 	if err != nil {
 		return nil, errors.Wrap(err, "query")
 	}
 
-	var segmentRanges []keyRange
+	var keyRanges []keyRange
 
-	err = cur.All(ctx, &segmentRanges)
+	err = cur.All(ctx, &keyRanges)
 	if err != nil {
 		return nil, errors.Wrap(err, "all")
 	}
 
-	return segmentRanges, nil
+	for i := range keyRanges {
+		log.Ctx(ctx).Debugf("Keyrange %d: type: %s, range [%v <=> %v]", i+1,
+			keyRanges[i].Min.Type.String(), keyRanges[i].Min, keyRanges[i].Max)
+	}
+
+	return keyRanges, nil
 }
 
 // CappedSegmenter provides sequential cursor access for capped collections.

tests/test_collections.py

@@ -632,13 +632,13 @@ def test_plm_110_rename_during_clone_and_repl(t: Testing):
     t.compare_all()
 
 
-def test_plm_126_clone_with_nan_id_document(t: Testing):
+def test_clone_with_nan_id_document(t: Testing):
     t.source["db_1"]["coll_1"].insert_one({"_id": float("nan"), "i": 100})
     t.source["db_1"]["coll_1"].insert_many(
         [{"_id": random.uniform(1e5, 1e10), "i": i} for i in range(50)]
     )
 
-    with t.run(phase=Runner.Phase.CLONE) as r:
+    with t.run(phase=Runner.Phase.MANUAL) as r:
         r.start()
         r.wait_for_clone_completed()
 
@@ -647,7 +647,6 @@ def test_plm_126_clone_with_nan_id_document(t: Testing):
     assert sourceDocCount == targetDocCount
 
 
-@pytest.mark.skip(reason="Clone with NaN _id is not supported for multi-id types")
 def test_clone_with_nan_id_document_multi_id_types(t: Testing):
     t.source["db_1"]["coll_1"].insert_one({"_id": Decimal128("NaN"), "i": 200})
     t.source["db_1"]["coll_1"].insert_many(
@@ -657,11 +656,13 @@ def test_clone_with_nan_id_document_multi_id_types(t: Testing):
         [{"_id": Decimal128(str(random.uniform(1e5, 1e10))), "i": i} for i in range(50)]
     )
     t.source["db_1"]["coll_1"].insert_many(
-        [{"_id": "inel" + str(random.uniform(1e5, 1e10)), "i": i} for i in range(50)]
+        [{"_id": str(random.uniform(1e5, 1e10)), "i": i} for i in range(50)]
     )
 
-    with t.run(phase=Runner.Phase.CLONE) as r:
+    with t.run(phase=Runner.Phase.MANUAL) as r:
         r.start()
         r.wait_for_clone_completed()
 
-    t.compare_all()
+    sourceDocCount = t.source["db_1"]["coll_1"].count_documents({})
+    targetDocCount = t.target["db_1"]["coll_1"].count_documents({})
+    assert sourceDocCount == targetDocCount