Add index hint support for distinct command (#1581)

Add index hint support for distinct command Add index hint support for distinct command for java, kotlin, and scala, and sync CRUD unified spec tests JAVA-5686
mongodb · Dec 12, 2024 · 08ec2fc · 08ec2fc
1 parent 6bb4005
commit 08ec2fc
Show file tree

Hide file tree

Showing 17 changed files with 330 additions and 9 deletions.
diff --git a/driver-core/src/main/com/mongodb/internal/operation/AsyncOperations.java b/driver-core/src/main/com/mongodb/internal/operation/AsyncOperations.java
@@ -181,8 +181,9 @@ public <TResult> AsyncReadOperation<AsyncBatchCursor<TResult>> find(final MongoN
     }
 
     public <TResult> AsyncReadOperation<AsyncBatchCursor<TResult>> distinct(final String fieldName, final Bson filter,
-            final Class<TResult> resultClass, final Collation collation, final BsonValue comment) {
-        return operations.distinct(fieldName, filter, resultClass, collation, comment);
+            final Class<TResult> resultClass, final Collation collation, final BsonValue comment, final Bson hint,
+            final String hintString) {
+        return operations.distinct(fieldName, filter, resultClass, collation, comment, hint, hintString);
     }
 
     public <TResult> AsyncExplainableReadOperation<AsyncBatchCursor<TResult>> aggregate(

diff --git a/driver-core/src/main/com/mongodb/internal/operation/DistinctOperation.java b/driver-core/src/main/com/mongodb/internal/operation/DistinctOperation.java
@@ -54,6 +54,7 @@ public class DistinctOperation<T> implements AsyncReadOperation<AsyncBatchCursor
     private BsonDocument filter;
     private Collation collation;
     private BsonValue comment;
+    private BsonValue hint;
 
     public DistinctOperation(final MongoNamespace namespace, final String fieldName, final Decoder<T> decoder) {
         this.namespace = notNull("namespace", namespace);
@@ -97,6 +98,15 @@ public DistinctOperation<T> comment(final BsonValue comment) {
         return this;
     }
 
+    public BsonValue getHint() {
+        return hint;
+    }
+
+    public DistinctOperation<T> hint(@Nullable final BsonValue hint) {
+        this.hint = hint;
+        return this;
+    }
+
     @Override
     public BatchCursor<T> execute(final ReadBinding binding) {
         return executeRetryableRead(binding, namespace.getDatabaseName(), getCommandCreator(), createCommandDecoder(),
@@ -124,6 +134,7 @@ private CommandCreator getCommandCreator() {
                 commandDocument.put("collation", collation.asDocument());
             }
             putIfNotNull(commandDocument, "comment", comment);
+            putIfNotNull(commandDocument, "hint", hint);
             return commandDocument;
         };
     }

diff --git a/driver-core/src/main/com/mongodb/internal/operation/Operations.java b/driver-core/src/main/com/mongodb/internal/operation/Operations.java
@@ -217,13 +217,20 @@ private <TResult> FindOperation<TResult> createFindOperation(final MongoNamespac
     }
 
     <TResult> DistinctOperation<TResult> distinct(final String fieldName, @Nullable final Bson filter, final Class<TResult> resultClass,
-            final Collation collation, final BsonValue comment) {
-        return new DistinctOperation<>(assertNotNull(namespace),
+            final Collation collation, final BsonValue comment, @Nullable final Bson hint, @Nullable final String hintString) {
+        DistinctOperation<TResult> operation = new DistinctOperation<>(assertNotNull(namespace),
                 fieldName, codecRegistry.get(resultClass))
                 .retryReads(retryReads)
                 .filter(filter == null ? null : filter.toBsonDocument(documentClass, codecRegistry))
                 .collation(collation)
                 .comment(comment);
+
+        if (hint != null) {
+            operation.hint(toBsonDocument(hint));
+        } else if (hintString != null) {
+            operation.hint(new BsonString(hintString));
+        }
+        return operation;
     }
 
     <TResult> AggregateOperation<TResult> aggregate(final List<? extends Bson> pipeline, final Class<TResult> resultClass,

diff --git a/driver-core/src/main/com/mongodb/internal/operation/SyncOperations.java b/driver-core/src/main/com/mongodb/internal/operation/SyncOperations.java
@@ -165,8 +165,9 @@ public <TResult> ReadOperation<BatchCursor<TResult>> find(final MongoNamespace f
 
     public <TResult> ReadOperation<BatchCursor<TResult>> distinct(final String fieldName, final Bson filter,
                                                                   final Class<TResult> resultClass,
-                                                                  final Collation collation, final BsonValue comment) {
-        return operations.distinct(fieldName, filter, resultClass, collation, comment);
+                                                                  final Collation collation, final BsonValue comment,
+                                                                  final Bson hint, final String hintString) {
+        return operations.distinct(fieldName, filter, resultClass, collation, comment, hint, hintString);
     }
 
     public <TResult> ExplainableReadOperation<BatchCursor<TResult>> aggregate(final List<? extends Bson> pipeline,

diff --git a/driver-core/src/test/resources/unified-test-format/crud/distinct-hint.json b/driver-core/src/test/resources/unified-test-format/crud/distinct-hint.json
@@ -0,0 +1,139 @@
+{
+  "description": "distinct-hint",
+  "schemaVersion": "1.0",
+  "runOnRequirements": [
+    {
+      "minServerVersion": "7.1.0"
+    }
+  ],
+  "createEntities": [
+    {
+      "client": {
+        "id": "client0",
+        "observeEvents": [
+          "commandStartedEvent"
+        ]
+      }
+    },
+    {
+      "database": {
+        "id": "database0",
+        "client": "client0",
+        "databaseName": "distinct-hint-tests"
+      }
+    },
+    {
+      "collection": {
+        "id": "collection0",
+        "database": "database0",
+        "collectionName": "coll0"
+      }
+    }
+  ],
+  "initialData": [
+    {
+      "collectionName": "coll0",
+      "databaseName": "distinct-hint-tests",
+      "documents": [
+        {
+          "_id": 1,
+          "x": 11
+        },
+        {
+          "_id": 2,
+          "x": 22
+        },
+        {
+          "_id": 3,
+          "x": 33
+        }
+      ]
+    }
+  ],
+  "tests": [
+    {
+      "description": "distinct with hint string",
+      "operations": [
+        {
+          "name": "distinct",
+          "object": "collection0",
+          "arguments": {
+            "fieldName": "x",
+            "filter": {
+              "_id": 1
+            },
+            "hint": "_id_"
+          },
+          "expectResult": [
+            11
+          ]
+        }
+      ],
+      "expectEvents": [
+        {
+          "client": "client0",
+          "events": [
+            {
+              "commandStartedEvent": {
+                "command": {
+                  "distinct": "coll0",
+                  "key": "x",
+                  "query": {
+                    "_id": 1
+                  },
+                  "hint": "_id_"
+                },
+                "commandName": "distinct",
+                "databaseName": "distinct-hint-tests"
+              }
+            }
+          ]
+        }
+      ]
+    },
+    {
+      "description": "distinct with hint document",
+      "operations": [
+        {
+          "name": "distinct",
+          "object": "collection0",
+          "arguments": {
+            "fieldName": "x",
+            "filter": {
+              "_id": 1
+            },
+            "hint": {
+              "_id": 1
+            }
+          },
+          "expectResult": [
+            11
+          ]
+        }
+      ],
+      "expectEvents": [
+        {
+          "client": "client0",
+          "events": [
+            {
+              "commandStartedEvent": {
+                "command": {
+                  "distinct": "coll0",
+                  "key": "x",
+                  "query": {
+                    "_id": 1
+                  },
+                  "hint": {
+                    "_id": 1
+                  }
+                },
+                "commandName": "distinct",
+                "databaseName": "distinct-hint-tests"
+              }
+            }
+          ]
+        }
+      ]
+    }
+  ]
+}
diff --git a/...ntegration/kotlin/com/mongodb/kotlin/client/coroutine/syncadapter/SyncDistinctIterable.kt b/...ntegration/kotlin/com/mongodb/kotlin/client/coroutine/syncadapter/SyncDistinctIterable.kt
@@ -33,6 +33,8 @@ data class SyncDistinctIterable<T : Any>(val wrapped: DistinctFlow<T>) :
     override fun collation(collation: Collation?): SyncDistinctIterable<T> = apply { wrapped.collation(collation) }
     override fun comment(comment: String?): SyncDistinctIterable<T> = apply { wrapped.comment(comment) }
     override fun comment(comment: BsonValue?): SyncDistinctIterable<T> = apply { wrapped.comment(comment) }
+    override fun hint(hint: Bson?): SyncDistinctIterable<T> = apply { wrapped.hint(hint) }
+    override fun hintString(hint: String?): SyncDistinctIterable<T> = apply { wrapped.hintString(hint) }
     override fun timeoutMode(timeoutMode: TimeoutMode): SyncDistinctIterable<T> = apply {
         wrapped.timeoutMode(timeoutMode)
     }

diff --git a/driver-kotlin-coroutine/src/main/kotlin/com/mongodb/kotlin/client/coroutine/DistinctFlow.kt b/driver-kotlin-coroutine/src/main/kotlin/com/mongodb/kotlin/client/coroutine/DistinctFlow.kt
@@ -19,6 +19,7 @@ import com.mongodb.annotations.Alpha
 import com.mongodb.annotations.Reason
 import com.mongodb.client.cursor.TimeoutMode
 import com.mongodb.client.model.Collation
+import com.mongodb.lang.Nullable
 import com.mongodb.reactivestreams.client.DistinctPublisher
 import java.util.concurrent.TimeUnit
 import kotlinx.coroutines.flow.Flow
@@ -103,5 +104,23 @@ public class DistinctFlow<T : Any>(private val wrapped: DistinctPublisher<T>) :
      */
     public fun comment(comment: BsonValue?): DistinctFlow<T> = apply { wrapped.comment(comment) }
 
+    /**
+     * Sets the hint for which index to use. A null value means no hint is set.
+     *
+     * @param hint the hint
+     * @return this
+     */
+    public fun hint(@Nullable hint: Bson?): DistinctFlow<T> = apply { wrapped.hint(hint) }
+
+    /**
+     * Sets the hint for which index to use. A null value means no hint is set.
+     *
+     * Note: If [DistinctFlow.hint] is set that will be used instead of any hint string.
+     *
+     * @param hint the name of the index which should be used for the operation
+     * @return this
+     */
+    public fun hintString(@Nullable hint: String?): DistinctFlow<T> = apply { wrapped.hintString(hint) }
+
     public override suspend fun collect(collector: FlowCollector<T>): Unit = wrapped.asFlow().collect(collector)
 }
diff --git a/...sync/src/integration/kotlin/com/mongodb/kotlin/client/syncadapter/SyncDistinctIterable.kt b/...sync/src/integration/kotlin/com/mongodb/kotlin/client/syncadapter/SyncDistinctIterable.kt
@@ -36,4 +36,6 @@ internal class SyncDistinctIterable<T : Any>(val wrapped: DistinctIterable<T>) :
     override fun collation(collation: Collation?): SyncDistinctIterable<T> = apply { wrapped.collation(collation) }
     override fun comment(comment: String?): SyncDistinctIterable<T> = apply { wrapped.comment(comment) }
     override fun comment(comment: BsonValue?): SyncDistinctIterable<T> = apply { wrapped.comment(comment) }
+    override fun hint(hint: Bson?): SyncDistinctIterable<T> = apply { wrapped.hint(hint) }
+    override fun hintString(hint: String?): SyncDistinctIterable<T> = apply { wrapped.hintString(hint) }
 }
diff --git a/driver-kotlin-sync/src/main/kotlin/com/mongodb/kotlin/client/DistinctIterable.kt b/driver-kotlin-sync/src/main/kotlin/com/mongodb/kotlin/client/DistinctIterable.kt
@@ -20,6 +20,7 @@ import com.mongodb.annotations.Reason
 import com.mongodb.client.DistinctIterable as JDistinctIterable
 import com.mongodb.client.cursor.TimeoutMode
 import com.mongodb.client.model.Collation
+import com.mongodb.lang.Nullable
 import java.util.concurrent.TimeUnit
 import org.bson.BsonValue
 import org.bson.conversions.Bson
@@ -99,4 +100,22 @@ public class DistinctIterable<T : Any?>(private val wrapped: JDistinctIterable<T
      * @return this
      */
     public fun comment(comment: BsonValue?): DistinctIterable<T> = apply { wrapped.comment(comment) }
+
+    /**
+     * Sets the hint for which index to use. A null value means no hint is set.
+     *
+     * @param hint the hint
+     * @return this
+     */
+    public fun hint(@Nullable hint: Bson?): DistinctIterable<T> = apply { wrapped.hint(hint) }
+
+    /**
+     * Sets the hint for which index to use. A null value means no hint is set.
+     *
+     * Note: If [DistinctFlow.hint] is set that will be used instead of any hint string.
+     *
+     * @param hint the name of the index which should be used for the operation
+     * @return this
+     */
+    public fun hintString(@Nullable hint: String?): DistinctIterable<T> = apply { wrapped.hintString(hint) }
 }
diff --git a/driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/DistinctPublisher.java b/driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/DistinctPublisher.java
@@ -97,6 +97,24 @@ public interface DistinctPublisher<TResult> extends Publisher<TResult> {
      */
     DistinctPublisher<TResult> comment(@Nullable BsonValue comment);
 
+    /**
+     * Sets the hint for which index to use. A null value means no hint is set.
+     *
+     * @param hint the hint
+     * @return this
+     * @since 5.3
+     */
+    DistinctPublisher<TResult> hint(@Nullable Bson hint);
+
+    /**
+     * Sets the hint for which index to use. A null value means no hint is set.
+     *
+     * @param hint the name of the index which should be used for the operation
+     * @return this
+     * @since 5.3
+     */
+    DistinctPublisher<TResult> hintString(@Nullable String hint);
+
     /**
      * Sets the timeoutMode for the cursor.
      *

diff --git a/...e-streams/src/main/com/mongodb/reactivestreams/client/internal/DistinctPublisherImpl.java b/...e-streams/src/main/com/mongodb/reactivestreams/client/internal/DistinctPublisherImpl.java
@@ -41,6 +41,8 @@ final class DistinctPublisherImpl<T> extends BatchCursorPublisher<T> implements
     private long maxTimeMS;
     private Collation collation;
     private BsonValue comment;
+    private Bson hint;
+    private String hintString;
 
     DistinctPublisherImpl(
             @Nullable final ClientSession clientSession,
@@ -88,6 +90,18 @@ public DistinctPublisher<T> comment(@Nullable final BsonValue comment) {
         return this;
     }
 
+    @Override
+    public DistinctPublisher<T> hint(@Nullable final Bson hint) {
+        this.hint = hint;
+        return this;
+    }
+
+    @Override
+    public DistinctPublisher<T> hintString(@Nullable final String hint) {
+        this.hintString = hint;
+        return this;
+    }
+
     @Override
     public DistinctPublisher<T> timeoutMode(final TimeoutMode timeoutMode) {
         super.timeoutMode(timeoutMode);
@@ -97,7 +111,7 @@ public DistinctPublisher<T> timeoutMode(final TimeoutMode timeoutMode) {
     @Override
     AsyncReadOperation<AsyncBatchCursor<T>> asAsyncReadOperation(final int initialBatchSize) {
         // initialBatchSize is ignored for distinct operations.
-        return getOperations().distinct(fieldName, filter, getDocumentClass(), collation, comment);
+        return getOperations().distinct(fieldName, filter, getDocumentClass(), collation, comment, hint, hintString);
     }
 
     @Override

diff --git a/.../test/functional/com/mongodb/reactivestreams/client/syncadapter/SyncDistinctIterable.java b/.../test/functional/com/mongodb/reactivestreams/client/syncadapter/SyncDistinctIterable.java
@@ -71,6 +71,18 @@ public DistinctIterable<T> comment(@Nullable final BsonValue comment) {
         return this;
     }
 
+    @Override
+    public DistinctIterable<T> hint(@Nullable final Bson hint) {
+        wrapped.hint(hint);
+        return this;
+    }
+
+    @Override
+    public DistinctIterable<T> hintString(@Nullable final String hint) {
+        wrapped.hintString(hint);
+        return this;
+    }
+
     @Override
     public DistinctIterable<T> timeoutMode(final TimeoutMode timeoutMode) {
         wrapped.timeoutMode(timeoutMode);

diff --git a/driver-scala/src/integration/scala/org/mongodb/scala/syncadapter/SyncDistinctIterable.scala b/driver-scala/src/integration/scala/org/mongodb/scala/syncadapter/SyncDistinctIterable.scala
@@ -61,4 +61,15 @@ case class SyncDistinctIterable[T](wrapped: DistinctObservable[T])
     wrapped.comment(comment)
     this
   }
+
+  override def hint(hint: Bson): DistinctIterable[T] = {
+    wrapped.hint(hint)
+    this
+  }
+
+  override def hintString(hint: String): DistinctIterable[T] = {
+    wrapped.hintString(hint)
+    this
+  }
+
 }