From 63dd07bdae59ed6a4003a7b27d2d9d50f3661668 Mon Sep 17 00:00:00 2001
From: jmarkerink <jmarkerink@outlook.com>
Date: Fri, 27 Feb 2026 23:02:19 +0100
Subject: [PATCH] feat: optimized aggregation pipeline to use indexes when
 possible

---
 .../mongo/backend/aggregation/Aggregation.java           | 9 +++++++++
 .../mongo/backend/aggregation/stage/MatchStage.java      | 4 ++++
 2 files changed, 13 insertions(+)

diff --git a/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/Aggregation.java b/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/Aggregation.java
index 206341f1..82658a72 100644
--- a/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/Aggregation.java
+++ b/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/Aggregation.java
@@ -189,6 +189,15 @@ public static Aggregation fromPipeline(List<Document> pipeline, DatabaseResolver
     }
 
     private List<Document> runStages() {
+        // Only apply index optimization if there are no variables to inject
+        // Variables need to be added to documents before $match can evaluate them
+        if (!hasVariables()) {
+            return stages.stream()
+                .filter(stage -> stage instanceof MatchStage)
+                .findFirst()
+                .map(match -> runStages(collection.handleQueryAsStream(((MatchStage) match).getQuery())))
+                .orElseGet(() -> runStages(collection.queryAllAsStream()));
+        }
         return runStages(collection.queryAllAsStream());
     }
 
diff --git a/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/MatchStage.java b/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/MatchStage.java
index 66efe63e..b11d4321 100644
--- a/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/MatchStage.java
+++ b/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/MatchStage.java
@@ -24,4 +24,8 @@ public String name() {
     public Stream<Document> apply(Stream<Document> stream) {
         return stream.filter(document -> queryMatcher.matches(document, query));
     }
+
+    public Document getQuery() {
+        return query;
+    }
 }