ExpressionIndexTest.java

import com.mongodb.client.MongoClients;
import com.mongodb.client.MongoCollection;
import org.bson.BsonArray;
import org.bson.BsonDocument;
import org.bson.BsonInt32;
import org.bson.BsonNull;
import org.bson.BsonString;
import org.bson.BsonValue;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.bson.json.JsonWriterSettings;

import java.util.ArrayList;
import java.util.List;

import static com.mongodb.client.model.Aggregates.match;
import static com.mongodb.client.model.Filters.and;
import static com.mongodb.client.model.Filters.eq;
import static com.mongodb.client.model.Filters.expr;
import static com.mongodb.client.model.Filters.gt;
import static com.mongodb.client.model.Filters.lt;
import static com.mongodb.client.model.Filters.nor;
import static com.mongodb.client.model.Filters.or;
import static com.mongodb.client.model.Indexes.ascending;

public class ExpressionIndexTest {
    public static void main(String[] args) throws InterruptedException {
        var size = 1_000_000;

        var client = MongoClients.create();

        var coll = client.getDatabase("test").getCollection("expressions");
        coll.drop();

        var list = new ArrayList<Document>(size);
        for (int i = 0; i < size; i++) {
            list.add(new Document("x", i));
        }
        coll.insertMany(list);
        coll.insertOne(new Document("x", null));

        coll.createIndex(ascending("x"));


        System.out.println("x = 5000 OR x = 6000");
        execute(coll,
                match(or(eq("x", 5000), eq("x", 6000))));
        execute(coll,
                match(expr(
                        new BsonDocument().append("$or",
                                new BsonArray(List.of(
                                        asComparisonExpression("$eq", "x", 5000),
                                        asComparisonExpression("$eq", "x", 6000)))))));
        System.out.println();

        System.out.println("x < 500 OR x > 999,500");
        execute(coll,
                match(or(lt("x", 500), gt("x", size - 500))));
        execute(coll,
                match(expr(
                        new BsonDocument().append("$or",
                                new BsonArray(List.of(
                                        asComparisonExpression("$lt", "x", 500),
                                        asComparisonExpression("$gt", "x", size - 500)))))));
        System.out.println();

        System.out.println("x > 5000 AND x < 6000");
        execute(coll,
                match(and(gt("x", 5000), lt("x", 6000))));
        execute(coll,
                match(expr(
                        new BsonDocument().append("$and",
                                new BsonArray(List.of(
                                        asComparisonExpression("$gt", "x", 5000),
                                        asComparisonExpression("$lt", "x", 6000)))))));

        System.out.println("x > 5000 AND not (x > 6000)");
        execute(coll,
                match(and(gt("x", 5000), nor(gt("x", 6000)))));
        execute(coll,
                match(expr(
                        new BsonDocument().append("$and",
                                new BsonArray(List.of(
                                        asComparisonExpression("$gt", "x", 5000),
                                        new BsonDocument().append("$not", new BsonArray(List.of(asComparisonExpression("$gt", "x", 6000))))))))));
        System.out.println();
    }

//    @SuppressWarnings("SameParameterValue")
//    private static BsonDocument asComparisonExpression(String operator, String fieldName, int fieldValue) {
//        return new BsonDocument().append(operator,
//                new BsonArray(List.of(new BsonString("$" + fieldName), new BsonInt32(fieldValue))));
//    }

    @SuppressWarnings("SameParameterValue")
    private static BsonDocument asComparisonExpression(String operator, String fieldName, int fieldValue) {
        var bsonFieldName = new BsonString("$" + fieldName);
        return new BsonDocument().append("$and",
                new BsonArray(List.of(
                        new BsonDocument().append("$gt",
                                new BsonArray(List.of(bsonFieldName, BsonNull.VALUE))),
                        new BsonDocument().append(operator,
                                new BsonArray(List.of(bsonFieldName, new BsonInt32(fieldValue)))))));
    }

    private static void execute(MongoCollection<Document> collection, Bson matchStage) {
        var aggregateIterable = collection.aggregate(List.of(matchStage));
        var settings = JsonWriterSettings.builder().indent(true).build();
        System.out.println(matchStage.toBsonDocument().toJson(settings));

        var plan = aggregateIterable.explain(BsonDocument.class);
        var bounds = getBounds(plan.getDocument("queryPlanner"));
        System.out.println("Index bounds: ");
        bounds.forEach(bound -> {
            System.out.print("   ");
            System.out.println(bound);});
        System.out.println();
        
        // warm server cache
        var numResults = aggregateIterable.into(new ArrayList<>()).size();
        System.out.println("Num results: " + numResults);

        long startTime = System.nanoTime();
        aggregateIterable.into(new ArrayList<>());
        long endTime = System.nanoTime();
        System.out.println("Elapsed (ms): " + (endTime - startTime) / 1_000_000.0);

        System.out.println();
    }

    private static List<BsonDocument> getBounds(BsonDocument plan) {
        var bounds = new ArrayList<BsonDocument>();
        plan.forEach((String key, BsonValue value) -> {
            if (key.equals("indexBounds")) {
                bounds.add(value.asDocument());
            }
            if (!key.equals("rejectedPlans")) {
                if (value.isDocument()) {
                    bounds.addAll(getBounds(value.asDocument()));
                } else if (value.isArray()) {
                    bounds.addAll(getBounds(value.asArray()));
                }
            }
        });
        return bounds;
    }

    private static List<BsonDocument> getBounds(BsonArray plan) {
        var bounds = new ArrayList<BsonDocument>();
        plan.forEach((BsonValue value) -> {
            if (value.isDocument()) {
                bounds.addAll(getBounds(value.asDocument()));
            } else if (value.isArray()) {
                bounds.addAll(getBounds(value.asArray()));
            }
        });
        return bounds;
    }
}