Add Sample

Author: clun

tools/src/main/java/com/datastax/astra/tool/CsvLoaderWiki.java

@@ -0,0 +1,66 @@
+package com.datastax.astra.tool;
+
+import com.datastax.astra.client.Collection;
+import com.datastax.astra.client.DataAPIClient;
+import com.datastax.astra.client.Database;
+import com.datastax.astra.client.model.CollectionOptions;
+import com.datastax.astra.client.model.Document;
+import com.datastax.astra.client.model.SimilarityMetric;
+import com.datastax.astra.tool.csv.CsvLoader;
+import com.datastax.astra.tool.csv.CsvRowMapper;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Load a CSV to Astra
+ */
+@Slf4j
+public class CsvLoaderWiki {
+
+    private static final String ASTRA_TOKEN  = "<CHANGE_ME>";
+    private static final String API_ENDPOINT = "<CHANGE_ME>";
+    private static final String CSV_FILE     = "<CHANGE_ME>";
+
+    public static void main(String[] args) throws Exception {
+        // Get an empty Collection
+        DataAPIClient client = new DataAPIClient(ASTRA_TOKEN);
+        Database wikiDataDb  = client.getDatabase(API_ENDPOINT);
+
+        Collection<Document> wiki = wikiDataDb.createCollection(
+                "wiki",
+                // Create collection with a Service in vectorize
+                CollectionOptions.builder()
+                        .vectorDimension(768) // found from the CSV
+                        .vectorSimilarity(SimilarityMetric.COSINE)
+                        .build());
+
+        // Optionally delete all documents
+        //wiki.deleteAll();
+
+        CsvLoader.load(CSV_FILE, wiki, new CsvRowMapper() {
+            int line = 0;
+            @Override
+            public Document map(Document doc) {
+                line++;
+                if (line % 1000 == 0) {
+                   log.info("{} lines have been submitted ",  line);
+                }
+                doc.put("_id", doc.get("", String.class));
+                doc.remove("");
+                String embedding = doc.get("embedding", String.class);
+                embedding = embedding.replaceAll("\\[", "");
+                embedding = embedding.replaceAll("\\]", "");
+                embedding = embedding.replaceAll("\\\n", "");
+                String[] parts = embedding.trim().split("\\s+");
+                float[] floatArray = new float[parts.length];
+                for (int i = 0; i < parts.length; i++) {
+                    floatArray[i] = Float.parseFloat(parts[i]);
+                }
+                doc.put("$vector", floatArray);
+                doc.remove("embedding");
+                return doc;
+            }
+        });
+
+    }
+
+}

tools/src/main/java/com/datastax/astra/tool/csv/CsvLoaderSettings.java

@@ -9,7 +9,7 @@ public class CsvLoaderSettings {
 
     private static final int THREAD_POOL_SIZE = 5;
 
-    private static final int TIMEOUT = 180;
+    private static final int TIMEOUT = 1800;
 
     @Builder.Default
     int batchSize = BATCH_SIZE;;