feat: shard key selector (#20)

Author: Anush008

README.md

@@ -213,6 +213,7 @@ Qdrant supports all the Spark data types. The appropriate types are mapped based
 | `sparse_vector_index_fields` | Comma-separated names of columns holding the sparse vector indices. | `ArrayType(IntegerType)`      | ❌       |
 | `sparse_vector_value_fields` | Comma-separated names of columns holding the sparse vector values.  | `ArrayType(FloatType)`        | ❌       |
 | `sparse_vector_names`        | Comma-separated names of the sparse vectors in the collection.      | -                             | ❌       |
+| `shard_key_selector`         | Comma-separated names of custom shard keys to use during upsert.    | -                             | ❌       |
 
 ## LICENSE 📜
 

src/main/java/io/qdrant/spark/QdrantDataWriter.java

@@ -81,7 +81,7 @@ public void write(int retries) {
     try {
       // Instantiate a new QdrantGrpc object to maintain serializability
       QdrantGrpc qdrant = new QdrantGrpc(new URL(this.qdrantUrl), this.apiKey);
-      qdrant.upsert(this.options.collectionName, this.points);
+      qdrant.upsert(this.options.collectionName, this.points, this.options.shardKeySelector);
       qdrant.close();
       this.points.clear();
     } catch (Exception e) {

src/main/java/io/qdrant/spark/QdrantGrpc.java

@@ -3,20 +3,25 @@
 import io.qdrant.client.QdrantClient;
 import io.qdrant.client.QdrantGrpcClient;
 import io.qdrant.client.grpc.Points.PointStruct;
+import io.qdrant.client.grpc.Points.ShardKeySelector;
+import io.qdrant.client.grpc.Points.UpsertPoints;
+
 import java.io.Serializable;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 
+import javax.annotation.Nullable;
+
 /** A class that provides methods to interact with Qdrant REST API. */
 public class QdrantGrpc implements Serializable {
   private final QdrantClient client;
 
   /**
    * Constructor for QdrantRest class.
    *
-   * @param url The URL of the Qdrant instance.
+   * @param url    The URL of the Qdrant instance.
    * @param apiKey The API key to authenticate with Qdrant.
    * @throws MalformedURLException If the URL is invalid.
    */
@@ -26,23 +31,33 @@ public QdrantGrpc(URL url, String apiKey) throws MalformedURLException {
     int port = url.getPort() == -1 ? 6334 : url.getPort();
     boolean useTls = url.getProtocol().equalsIgnoreCase("https");
 
-    this.client =
-        new QdrantClient(
-            QdrantGrpcClient.newBuilder(host, port, useTls).withApiKey(apiKey).build());
+    this.client = new QdrantClient(
+        QdrantGrpcClient.newBuilder(host, port, useTls).withApiKey(apiKey).build());
   }
 
   /**
    * Uploads a batch of points to a Qdrant collection.
    *
    * @param collectionName The name of the collection to upload the points to.
-   * @param points The list of points to upload.
-   * @throws InterruptedException If there was an error uploading the batch to Qdrant.
-   * @throws ExecutionException If there was an error uploading the batch to Qdrant.
+   * @param points         The list of points to upload.
+   * @param shardKeySelector The shard key selector to use for the upsert.
+   * 
+   * @throws InterruptedException If there was an error uploading the batch to
+   *                              Qdrant.
+   * @throws ExecutionException   If there was an error uploading the batch to
+   *                              Qdrant.
    */
-  public void upsert(String collectionName, List<PointStruct> points)
+  public void upsert(String collectionName, List<PointStruct> points, @Nullable ShardKeySelector shardKeySelector)
       throws InterruptedException, ExecutionException {
-    this.client.upsertAsync(collectionName, points).get();
-    return;
+
+    UpsertPoints.Builder upsertPoints = UpsertPoints.newBuilder().setCollectionName(collectionName)
+        .addAllPoints(points);
+
+    if (shardKeySelector != null) {
+      upsertPoints.setShardKeySelector(shardKeySelector);
+    }
+
+    this.client.upsertAsync(upsertPoints.build()).get();
   }
 
   public void close() {

src/main/java/io/qdrant/spark/QdrantOptions.java

@@ -7,6 +7,14 @@
 import java.util.Map;
 import java.util.Objects;
 
+import javax.annotation.Nullable;
+
+import io.qdrant.client.grpc.Collections.ShardKey;
+import io.qdrant.client.grpc.Points.ShardKeySelector;
+
+import static io.qdrant.client.ShardKeySelectorFactory.shardKeySelector;
+import static io.qdrant.client.ShardKeyFactory.shardKey;
+
 public class QdrantOptions implements Serializable {
   private static final int DEFAULT_BATCH_SIZE = 64;
   private static final int DEFAULT_RETRIES = 3;
@@ -25,14 +33,14 @@ public class QdrantOptions implements Serializable {
   public final String[] vectorFields;
   public final String[] vectorNames;
   public final List<String> payloadFieldsToSkip;
+  public final ShardKeySelector shardKeySelector;
 
   public QdrantOptions(Map<String, String> options) {
     Objects.requireNonNull(options);
 
     qdrantUrl = options.get("qdrant_url");
     collectionName = options.get("collection_name");
-    batchSize =
-        Integer.parseInt(options.getOrDefault("batch_size", String.valueOf(DEFAULT_BATCH_SIZE)));
+    batchSize = Integer.parseInt(options.getOrDefault("batch_size", String.valueOf(DEFAULT_BATCH_SIZE)));
     retries = Integer.parseInt(options.getOrDefault("retries", String.valueOf(DEFAULT_RETRIES)));
     idField = options.getOrDefault("id_field", "");
     apiKey = options.getOrDefault("api_key", "");
@@ -45,6 +53,8 @@ public QdrantOptions(Map<String, String> options) {
     vectorFields = parseArray(options.get("vector_fields"));
     vectorNames = parseArray(options.get("vector_names"));
 
+    shardKeySelector = parseShardKeys(options.get("shard_key_selector"));
+
     validateSparseVectorFields();
     validateVectorFields();
 
@@ -83,4 +93,33 @@ private void validateVectorFields() {
       throw new IllegalArgumentException("Vector fields and names should have the same length");
     }
   }
+
+  private ShardKeySelector parseShardKeys(@Nullable String shardKeys) {
+    if (shardKeys == null) {
+      return null;
+    }
+    String[] keys = shardKeys.split(",");
+
+    ShardKey[] shardKeysArray = new ShardKey[keys.length];
+
+    for (int i = 0; i < keys.length; i++) {
+      String key = keys[i];
+      if (isInt(key.trim())) {
+        shardKeysArray[i] = shardKey(Integer.parseInt(key.trim()));
+      } else {
+        shardKeysArray[i] = shardKey(key.trim());
+      }
+    }
+
+    return shardKeySelector(shardKeysArray);
+  }
+
+  boolean isInt(String s) {
+    try {
+      Integer.parseInt(s);
+      return true;
+    } catch (NumberFormatException er) {
+      return false;
+    }
+  }
 }

src/test/java/io/qdrant/spark/TestQdrantGrpc.java

@@ -69,7 +69,7 @@ public void testUploadBatch() throws Exception {
     points.add(point2Builder.build());
 
     // call the uploadBatch method
-    qdrantGrpc.upsert(collectionName, points);
+    qdrantGrpc.upsert(collectionName, points, null);
 
     qdrantGrpc.close();
   }