From e152ecc2f0d519be0cacc93207c9a71cab4b919e Mon Sep 17 00:00:00 2001 From: Eiji Ogiwara <146042372+ogiogidayo@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:20:24 +0900 Subject: [PATCH] Fix MongoDbIO read splitting to preserve non-ObjectId _id types splitKeysToFilters formatted split boundaries unconditionally as ObjectId("..."), so collections keyed by non-ObjectId ids (e.g. application-defined string ids) failed to parse the generated range filters when splitting was enabled via bucketAuto or numSplits. Serialize boundaries with their actual BSON types instead, and keep the raw id type in splitKeysToMatch as well. Fixes #39900 --- CHANGES.md | 1 + .../apache/beam/sdk/io/mongodb/MongoDbIO.java | 69 ++++++++++--------- .../beam/sdk/io/mongodb/MongoDbIOTest.java | 52 +++++++++++--- 3 files changed, 81 insertions(+), 41 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 151cbe14aaa2..05e4a7ea42e0 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -89,6 +89,7 @@ * (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)). * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)). * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)). +* (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)). ## Security Fixes diff --git a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java index 46c3f8fcd58a..501eaa4867c4 100644 --- a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java +++ b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java @@ -67,11 +67,9 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.bson.BsonDocument; import org.bson.BsonInt32; -import org.bson.BsonObjectId; import org.bson.BsonString; import org.bson.Document; import org.bson.conversions.Bson; -import org.bson.types.ObjectId; import org.checkerframework.checker.nullness.qual.Nullable; import org.checkerframework.dataflow.qual.Pure; import org.slf4j.Logger; @@ -606,38 +604,26 @@ public List> split( @VisibleForTesting static List splitKeysToFilters(List splitKeys) { ArrayList filters = new ArrayList<>(); - String lowestBound = null; // lower boundary (previous split in the iteration) + Object lowestBound = null; // lower boundary (previous split in the iteration) for (int i = 0; i < splitKeys.size(); i++) { - String splitKey = splitKeys.get(i).get("_id").toString(); - String rangeFilter; + Object splitKey = splitKeys.get(i).get("_id"); if (i == 0) { // this is the first split in the list, the filter defines // the range from the beginning up to this split - rangeFilter = String.format("{ $and: [ {\"_id\":{$lte:ObjectId(\"%s\")}}", splitKey); - filters.add(String.format("%s ]}", rangeFilter)); + filters.add(rangeFilter(null, splitKey)); // If there is only one split, also generate a range from the split to the end if (splitKeys.size() == 1) { - rangeFilter = String.format("{ $and: [ {\"_id\":{$gt:ObjectId(\"%s\")}}", splitKey); - filters.add(String.format("%s ]}", rangeFilter)); + filters.add(rangeFilter(splitKey, null)); } } else if (i == splitKeys.size() - 1) { // this is the last split in the list, the filters define // the range from the previous split to the current split and also // the current split to the end - rangeFilter = - String.format( - "{ $and: [ {\"_id\":{$gt:ObjectId(\"%s\")," + "$lte:ObjectId(\"%s\")}}", - lowestBound, splitKey); - filters.add(String.format("%s ]}", rangeFilter)); - rangeFilter = String.format("{ $and: [ {\"_id\":{$gt:ObjectId(\"%s\")}}", splitKey); - filters.add(String.format("%s ]}", rangeFilter)); + filters.add(rangeFilter(lowestBound, splitKey)); + filters.add(rangeFilter(splitKey, null)); } else { // we are between two splits - rangeFilter = - String.format( - "{ $and: [ {\"_id\":{$gt:ObjectId(\"%s\")," + "$lte:ObjectId(\"%s\")}}", - lowestBound, splitKey); - filters.add(String.format("%s ]}", rangeFilter)); + filters.add(rangeFilter(lowestBound, splitKey)); } lowestBound = splitKey; @@ -646,6 +632,23 @@ static List splitKeysToFilters(List splitKeys) { return filters; } + /** + * Builds a JSON range filter on {@code _id} with the given bounds. Bounds are serialized with + * their actual BSON types (as extended JSON) so that ids that are not ObjectIds, such as + * application-defined string ids, are preserved. + */ + private static String rangeFilter( + @Nullable Object greaterThan, @Nullable Object lessThanOrEqualTo) { + Document range = new Document(); + if (greaterThan != null) { + range.append("$gt", greaterThan); + } + if (lessThanOrEqualTo != null) { + range.append("$lte", lessThanOrEqualTo); + } + return new Document("$and", Collections.singletonList(new Document("_id", range))).toJson(); + } + /** * Transform a list of split keys as a list of filters containing corresponding range. * @@ -674,9 +677,11 @@ static List splitKeysToFilters(List splitKeys) { @VisibleForTesting static List splitKeysToMatch(List splitKeys) { List aggregates = new ArrayList<>(); - ObjectId lowestBound = null; // lower boundary (previous split in the iteration) + Object lowestBound = null; // lower boundary (previous split in the iteration) for (int i = 0; i < splitKeys.size(); i++) { - ObjectId splitKey = splitKeys.get(i).getObjectId("_id"); + // Keep the raw value so that ids that are not ObjectIds, such as application-defined + // string ids, are preserved. + Object splitKey = splitKeys.get(i).get("_id"); if (i == 0) { aggregates.add(Aggregates.match(Filters.lte("_id", splitKey))); if (splitKeys.size() == 1) { @@ -687,22 +692,20 @@ static List splitKeysToMatch(List splitKeys) { // the range from the previous split to the current split and also // the current split to the end // Create a custom BSON document with multiple conditions on the same field - BsonDocument rangeFilter = - new BsonDocument( + Document rangeFilter = + new Document( "_id", - new BsonDocument( - "$gt", new BsonObjectId(Preconditions.checkStateNotNull(lowestBound))) - .append("$lte", new BsonObjectId(splitKey))); + new Document("$gt", Preconditions.checkStateNotNull(lowestBound)) + .append("$lte", splitKey)); aggregates.add(Aggregates.match(rangeFilter)); aggregates.add(Aggregates.match(Filters.gt("_id", splitKey))); } else { // Create a custom BSON document with multiple conditions on the same field - BsonDocument rangeFilter = - new BsonDocument( + Document rangeFilter = + new Document( "_id", - new BsonDocument( - "$gt", new BsonObjectId(Preconditions.checkStateNotNull(lowestBound))) - .append("$lte", new BsonObjectId(splitKey))); + new Document("$gt", Preconditions.checkStateNotNull(lowestBound)) + .append("$lte", splitKey)); aggregates.add(Aggregates.match(rangeFilter)); } diff --git a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java index 94b9df527d2b..e0a0eabd5be3 100644 --- a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java +++ b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java @@ -103,20 +103,40 @@ public void testSplitIntoFilters() { documents.add(new Document("_id", 56)); List filters = MongoDbIO.BoundedMongoDbSource.splitKeysToFilters(documents); assertEquals(2, filters.size()); - assertEquals("{ $and: [ {\"_id\":{$lte:ObjectId(\"56\")}} ]}", filters.get(0)); - assertEquals("{ $and: [ {\"_id\":{$gt:ObjectId(\"56\")}} ]}", filters.get(1)); + assertEquals(56, idRange(filters.get(0)).getInt32("$lte").getValue()); + assertEquals(56, idRange(filters.get(1)).getInt32("$gt").getValue()); // Add two more splits; now we should have 4 filters documents.add(new Document("_id", 109)); documents.add(new Document("_id", 256)); filters = MongoDbIO.BoundedMongoDbSource.splitKeysToFilters(documents); assertEquals(4, filters.size()); - assertEquals("{ $and: [ {\"_id\":{$lte:ObjectId(\"56\")}} ]}", filters.get(0)); - assertEquals( - "{ $and: [ {\"_id\":{$gt:ObjectId(\"56\"),$lte:ObjectId(\"109\")}} ]}", filters.get(1)); - assertEquals( - "{ $and: [ {\"_id\":{$gt:ObjectId(\"109\"),$lte:ObjectId(\"256\")}} ]}", filters.get(2)); - assertEquals("{ $and: [ {\"_id\":{$gt:ObjectId(\"256\")}} ]}", filters.get(3)); + assertEquals(56, idRange(filters.get(0)).getInt32("$lte").getValue()); + assertEquals(56, idRange(filters.get(1)).getInt32("$gt").getValue()); + assertEquals(109, idRange(filters.get(1)).getInt32("$lte").getValue()); + assertEquals(109, idRange(filters.get(2)).getInt32("$gt").getValue()); + assertEquals(256, idRange(filters.get(2)).getInt32("$lte").getValue()); + assertEquals(256, idRange(filters.get(3)).getInt32("$gt").getValue()); + } + + @Test + public void testSplitIntoFiltersWithStringId() { + // Ids that are not ObjectIds, such as application-defined string ids, must keep their type: + // they used to be formatted as ObjectId("...") and could never be parsed back (#39900). + ArrayList documents = new ArrayList<>(); + documents.add(new Document("_id", "id-aaa")); + documents.add(new Document("_id", "id-mmm")); + List filters = MongoDbIO.BoundedMongoDbSource.splitKeysToFilters(documents); + assertEquals(3, filters.size()); + assertEquals("id-aaa", idRange(filters.get(0)).getString("$lte").getValue()); + assertEquals("id-aaa", idRange(filters.get(1)).getString("$gt").getValue()); + assertEquals("id-mmm", idRange(filters.get(1)).getString("$lte").getValue()); + assertEquals("id-mmm", idRange(filters.get(2)).getString("$gt").getValue()); + } + + /** Parses a filter generated by splitKeysToFilters and returns the range on {@code _id}. */ + private static BsonDocument idRange(String filter) { + return BsonDocument.parse(filter).getArray("$and").get(0).asDocument().getDocument("_id"); } @Test @@ -152,6 +172,22 @@ public void testSplitIntoBucket() { buckets.get(3).toString()); } + @Test + public void testSplitIntoBucketWithStringId() { + // Ids that are not ObjectIds, such as application-defined string ids, must keep their type + // instead of being read as ObjectIds (#39900). + ArrayList documents = new ArrayList<>(); + documents.add(new Document("_id", "id-aaa")); + documents.add(new Document("_id", "id-mmm")); + List buckets = MongoDbIO.BoundedMongoDbSource.splitKeysToMatch(documents); + assertEquals(3, buckets.size()); + assertEquals("{\"$match\": {\"_id\": {\"$lte\": \"id-aaa\"}}}", buckets.get(0).toString()); + assertEquals( + "{\"$match\": {\"_id\": {\"$gt\": \"id-aaa\", \"$lte\": \"id-mmm\"}}}", + buckets.get(1).toString()); + assertEquals("{\"$match\": {\"_id\": {\"$gt\": \"id-mmm\"}}}", buckets.get(2).toString()); + } + @Test public void testBuildAutoBuckets() { List aggregates = new ArrayList();