Added map constructor for PGsparsevec - #10

Co-authored-by: tkwlsrl <[email protected]>
pgvector · Jun 20, 2024 · c8ea95d · c8ea95d
1 parent 6352966
commit c8ea95d
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 0 deletions.
diff --git a/src/main/java/com/pgvector/PGsparsevec.java b/src/main/java/com/pgvector/PGsparsevec.java
@@ -4,7 +4,9 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.Arrays;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import org.postgresql.PGConnection;
 import org.postgresql.util.ByteConverter;
@@ -92,6 +94,36 @@ public <T extends Number> PGsparsevec(List<T> v) {
         }
     }
 
+    /**
+     * Constructor
+     *
+     * @param map <Integer, T> map of non-zero elements
+     * @param dimensions number of dimensions
+     */
+    public <T extends Number> PGsparsevec(Map<Integer, T> map, int dimensions) {
+        this();
+
+        ArrayList<Map.Entry<Integer, T>> elements = new ArrayList<Map.Entry<Integer, T>>();
+        if (!Objects.isNull(map)) {
+            elements.addAll(map.entrySet());
+        }
+        elements.removeIf((e) -> e.getValue().floatValue() == 0);
+        elements.sort((a, b) -> Integer.compare(a.getKey(), b.getKey()));
+
+        int nnz = elements.size();
+        indices = new int[nnz];
+        values = new float[nnz];
+
+        int i = 0;
+        for (Map.Entry<Integer, T> e : elements) {
+            indices[i] = e.getKey().intValue();
+            values[i] = e.getValue().floatValue();
+            i++;
+        }
+
+        this.dimensions = dimensions;
+    }
+
     /**
      * Constructor
      *

diff --git a/src/test/java/com/pgvector/PGsparsevecTest.java b/src/test/java/com/pgvector/PGsparsevecTest.java
@@ -2,6 +2,8 @@
 
 import java.sql.SQLException;
 import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
 import com.pgvector.PGsparsevec;
 import org.junit.jupiter.api.Test;
 
@@ -35,6 +37,18 @@ void testDoubleListConstructor() {
         assertArrayEquals(new float[] {1, 2, 3}, vec.toArray());
     }
 
+    @Test
+    void testMapConstructor() {
+        Map<Integer, Float> map = new HashMap<Integer, Float>();
+        map.put(Integer.valueOf(2), Float.valueOf(2));
+        map.put(Integer.valueOf(4), Float.valueOf(3));
+        map.put(Integer.valueOf(0), Float.valueOf(1));
+        map.put(Integer.valueOf(3), Float.valueOf(0));
+        PGsparsevec vec = new PGsparsevec(map, 6);
+        assertArrayEquals(new float[] {1, 0, 2, 0, 3, 0}, vec.toArray());
+        assertArrayEquals(new int[] {0, 2, 4}, vec.getIndices());
+    }
+
     @Test
     void testGetValue() {
         PGsparsevec vec = new PGsparsevec(new float[] {1, 0, 2, 0, 3, 0});