Enable setting custom vertex ids in TP3 interface #147

janusgraph-core/src/main/java/org/janusgraph/core/JanusGraphTransaction.java

@@ -49,9 +49,13 @@ public interface JanusGraphTransaction extends Transaction {
      * Note, that an exception is thrown if the vertex id is not a valid JanusGraph vertex id or if a vertex with the given
      * id already exists.
      * <p/>
-     * Custom id setting must be enabled via the configuration option {@link org.janusgraph.graphdb.configuration.GraphDatabaseConfiguration#ALLOW_SETTING_VERTEX_ID}.
-     * <p/>
-     * Use {@link org.janusgraph.core.util.JanusGraphId#toVertexId(long)} to construct a valid JanusGraph vertex id from a user id.
+     * Custom id setting must be enabled via the configuration option {@link org.janusgraph.graphdb.configuration.GraphDatabaseConfiguration#ALLOW_SETTING_VERTEX_ID}
+     * and valid JanusGraph vertex ids must be provided. Use {@link org.janusgraph.graphdb.idmanagement.IDManager#toVertexId(long)}
+     * to construct a valid JanusGraph vertex id from a user id, where <code>idManager</code> can be obtained through
+     * {@link org.janusgraph.graphdb.database.StandardJanusGraph#getIDManager()}.
+     * <pre>
+     * <code>long vertexId = ((StandardJanusGraph) graph).getIDManager().toVertexId(userVertexId);</code>
+     * </pre>
      *
      * @param id vertex id of the vertex to be created
      * @param vertexLabel vertex label for this vertex - can be null if no vertex label should be set.

janusgraph-core/src/main/java/org/janusgraph/core/util/JanusGraphId.java

@@ -22,7 +22,17 @@
  * Utility methods for handling JanusGraph ids and converting them between indexing and storage backend representations.
  *
  * @author Matthias Broecheler ([email protected])
+ * @deprecated This class does not produce valid JanusGraph vertex ids as it does not take into account partitioning
+ * bits used in vertex id assignment. Use {@link org.janusgraph.graphdb.idmanagement.IDManager}, which can be obtained
+ * through {@link org.janusgraph.graphdb.database.StandardJanusGraph#getIDManager()} and includes methods for converting
+ * a user id to ({@link org.janusgraph.graphdb.idmanagement.IDManager#toVertexId(long)}) and from
+ * ({@link org.janusgraph.graphdb.idmanagement.IDManager#fromVertexId(long)}) JanusGraph vertex id.
+ * <p/>
+ * <pre>
+ * <code>IDManager idManager = ((StandardJanusGraph) graph).getIDManager();</code>
+ * </pre>
  */
+@Deprecated
 public class JanusGraphId {
 
     /**
@@ -31,6 +41,7 @@ public class JanusGraphId {
      *
      * @param id long id
      * @return a corresponding JanusGraph vertex id
+     * @deprecated Use {@link org.janusgraph.graphdb.idmanagement.IDManager#toVertexId(long)}.
      */
     public static final long toVertexId(long id) {
         Preconditions.checkArgument(id > 0, "Vertex id must be positive: %s", id);
@@ -43,6 +54,7 @@ public static final long toVertexId(long id) {
      *
      * @param id JanusGraph vertex id (must be positive)
      * @return original user provided id
+     * @deprecated Use {@link org.janusgraph.graphdb.idmanagement.IDManager#fromVertexId(long)}
      */
     public static final long fromVertexId(long id) {
         Preconditions.checkArgument(id > 0, "Invalid vertex id provided: %s", id);
@@ -54,6 +66,7 @@ public static final long fromVertexId(long id) {
      *
      * @param v Vertex
      * @return original user provided id
+     * @deprecated Use {@link org.janusgraph.graphdb.idmanagement.IDManager#fromVertexId(long)}
      */
     public static final long fromVertexID(JanusGraphVertex v) {
         Preconditions.checkArgument(v.hasId(), "Invalid vertex provided: %s", v);

janusgraph-core/src/main/java/org/janusgraph/graphdb/idmanagement/IDManager.java

@@ -567,6 +567,33 @@ public long[] getPartitionedVertexRepresentatives(long partitionedVertexId) {
         return ids;
     }
 
+    /**
+     * Converts a user provided long id into a JanusGraph vertex id. The id must be positive and less than {@link #getVertexCountBound()}.
+     * This method is useful when providing ids during vertex creation via {@link org.apache.tinkerpop.gremlin.structure.Graph#addVertex(Object...)}.
+     *
+     * @param id long id
+     * @return a corresponding JanusGraph vertex id
+     * @see #fromVertexId(long)
+     */
+    public long toVertexId(long id) {
+        Preconditions.checkArgument(id > 0, "Vertex id must be positive: %s", id);
+        Preconditions.checkArgument(vertexCountBound > id, "Vertex id is too large: %s", id);
+        return id<<(partitionBits+USERVERTEX_PADDING_BITWIDTH);
+    }
+
+    /**
+     * Converts a JanusGraph vertex id to the user provided id as the inverse mapping of {@link #toVertexId(long)}.
+     *
+     * @param id JanusGraph vertex id (must be positive)
+     * @return original user provided id
+     * @see #toVertexId(long)
+     */
+    public long fromVertexId(long id) {
+        Preconditions.checkArgument(id >>> USERVERTEX_PADDING_BITWIDTH+partitionBits > 0
+            && id <= (vertexCountBound-1)<<USERVERTEX_PADDING_BITWIDTH+partitionBits, "Invalid vertex id provided: %s", id);
+        return id>>USERVERTEX_PADDING_BITWIDTH+partitionBits;
+    }
+
     public boolean isPartitionedVertex(long id) {
         return isUserVertexId(id) && VertexIDType.PartitionedVertex.is(id);
     }

janusgraph-core/src/main/java/org/janusgraph/graphdb/tinkerpop/JanusGraphBlueprintsTransaction.java

@@ -19,6 +19,7 @@
 import org.janusgraph.core.JanusGraphVertex;
 import org.janusgraph.core.VertexLabel;
 import org.janusgraph.diskstorage.util.Hex;
+import org.janusgraph.graphdb.database.StandardJanusGraph;
 import org.janusgraph.graphdb.olap.computer.FulgoraGraphComputer;
 import org.janusgraph.graphdb.relations.RelationIdentifier;
 import org.janusgraph.graphdb.types.system.BaseVertexLabel;
@@ -89,17 +90,21 @@ public FulgoraGraphComputer compute() throws IllegalArgumentException {
      * Note, that an exception is thrown if the vertex id is not a valid JanusGraph vertex id or if a vertex with the given
      * id already exists. Only accepts long ids - all others are ignored.
      * <p/>
-     * Custom id setting must be enabled via the configuration option {@link org.janusgraph.graphdb.configuration.GraphDatabaseConfiguration#ALLOW_SETTING_VERTEX_ID}.
-     * <p/>
-     * Use {@link org.janusgraph.core.util.JanusGraphId#toVertexId(long)} to construct a valid JanusGraph vertex id from a user id.
+     * Custom id setting must be enabled via the configuration option {@link org.janusgraph.graphdb.configuration.GraphDatabaseConfiguration#ALLOW_SETTING_VERTEX_ID}
+     * and valid JanusGraph vertex ids must be provided. Use {@link org.janusgraph.graphdb.idmanagement.IDManager#toVertexId(long)}
+     * to construct a valid JanusGraph vertex id from a user id, where <code>idManager</code> can be obtained through
+     * {@link org.janusgraph.graphdb.database.StandardJanusGraph#getIDManager()}.
+     * <pre>
+     * <code>long vertexId = ((StandardJanusGraph) graph).getIDManager().toVertexId(userVertexId);</code>
+     * </pre>
      *
      * @param keyValues key-value pairs of properties to characterize or attach to the vertex
      * @return New vertex
      */
     @Override
     public JanusGraphVertex addVertex(Object... keyValues) {
         ElementHelper.legalPropertyKeyValueArray(keyValues);
-        if (ElementHelper.getIdValue(keyValues).isPresent()) throw Vertex.Exceptions.userSuppliedIdsNotSupported();
+        if (ElementHelper.getIdValue(keyValues).isPresent() && !((StandardJanusGraph) getGraph()).getConfiguration().allowVertexIdSetting()) throw Vertex.Exceptions.userSuppliedIdsNotSupported();
         Object labelValue = null;
         for (int i = 0; i < keyValues.length; i = i + 2) {
             if (keyValues[i].equals(T.label)) {
@@ -114,11 +119,8 @@ public JanusGraphVertex addVertex(Object... keyValues) {
             label = (labelValue instanceof VertexLabel)?(VertexLabel)labelValue:getOrCreateVertexLabel((String) labelValue);
         }
 
-        final JanusGraphVertex vertex = addVertex(null,label);
-//        for (int i = 0; i < keyValues.length; i = i + 2) {
-//            if (!keyValues[i].equals(T.id) && !keyValues[i].equals(T.label))
-//                ((StandardJanusGraphTx)this).addPropertyInternal(vertex,getOrCreatePropertyKey((String) keyValues[i]),keyValues[i+1]);
-//        }
+        final Long id = ElementHelper.getIdValue(keyValues).map(Number.class::cast).map(Number::longValue).orElse(null);
+        final JanusGraphVertex vertex = addVertex(id, label);
         org.janusgraph.graphdb.util.ElementHelper.attachProperties(vertex, keyValues);
         return vertex;
     }

janusgraph-test/src/test/java/org/janusgraph/graphdb/idmanagement/VertexIDAssignerTest.java

@@ -16,6 +16,7 @@
 
 import com.carrotsearch.hppc.LongHashSet;
 import com.carrotsearch.hppc.LongSet;
+import org.apache.tinkerpop.gremlin.structure.T;
 import org.janusgraph.core.JanusGraphFactory;
 import org.janusgraph.core.JanusGraph;
 import org.janusgraph.core.JanusGraphVertex;
@@ -68,6 +69,12 @@ public static Collection<Object[]> configs() {
 
     private final long maxIDAssignments;
 
+    private final long numPartitionsBits;
+
+    private enum CustomIdStrategy {
+        LOW,
+        HIGH
+    }
 
     /**
      *
@@ -97,13 +104,17 @@ public VertexIDAssignerTest(int numPartitionsBits, int partitionMax, int[] local
         } else {
             this.maxIDAssignments = (1<<numPartitionsBits)*((long)partitionMax);
         }
+
+        this.numPartitionsBits = numPartitionsBits;
     }
 
-    private static JanusGraph getInMemoryGraph() {
+    private JanusGraph getInMemoryGraph(boolean allowSettingVertexId, boolean idsFlush) {
         ModifiableConfiguration config = GraphDatabaseConfiguration.buildGraphConfiguration();
         config.set(GraphDatabaseConfiguration.STORAGE_BACKEND, InMemoryStoreManager.class.getCanonicalName());
-        config.set(GraphDatabaseConfiguration.IDS_FLUSH, false);
+        config.set(GraphDatabaseConfiguration.IDS_FLUSH, idsFlush);
         config.set(GraphDatabaseConfiguration.IDAUTHORITY_WAIT, Duration.ofMillis(1L));
+        config.set(GraphDatabaseConfiguration.CLUSTER_MAX_PARTITIONS, 1<<numPartitionsBits);
+        config.set(GraphDatabaseConfiguration.ALLOW_SETTING_VERTEX_ID, allowSettingVertexId);
         return JanusGraphFactory.open(config);
     }
 
@@ -115,7 +126,7 @@ public void testIDAssignment() {
         int totalVertices = 0;
         for (int trial = 0; trial < 10; trial++) {
             for (boolean flush : new boolean[]{true, false}) {
-                JanusGraph graph = getInMemoryGraph();
+                JanusGraph graph = getInMemoryGraph(false, false);
                 int numVertices = 1000;
                 List<JanusGraphVertex> vertices = new ArrayList<JanusGraphVertex>(numVertices);
                 List<InternalRelation> relations = new ArrayList<InternalRelation>();
@@ -171,5 +182,70 @@ public void testIDAssignment() {
         }
     }
 
+    @Test
+    public void testCustomIdAssignment() {
+        testCustomIdAssignment(CustomIdStrategy.LOW);
+        testCustomIdAssignment(CustomIdStrategy.HIGH);
+
+        final IDManager idManager = idAssigner.getIDManager();
+        for (final long id : new long[] {0, idManager.getVertexCountBound()}) {
+            try {
+                idManager.toVertexId(id);
+                fail("Should fail to convert out of range user id to graph vertex id");
+            } catch (IllegalArgumentException e) {
+                // should throw this exception
+            }
+        }
+
+        for (final long vertexId : new long[] {idManager.toVertexId(1)-1, idManager.toVertexId(idManager.getVertexCountBound()-1)+1}) {
+            try {
+                idManager.fromVertexId(vertexId);
+                fail("Should fail to convert out of range vertex id to user id");
+            } catch (IllegalArgumentException e) {
+                // should throw this exception
+            }
+        }
+    }
+
+    private void testCustomIdAssignment(CustomIdStrategy idStrategy) {
+        LongSet vertexIds = new LongHashSet();
+        final long maxCount = idAssigner.getIDManager().getVertexCountBound();
+        long count = 1;
+        for (int trial = 0; trial < 10; trial++) {
+            JanusGraph graph = getInMemoryGraph(true, true);
+            int numVertices = 1000;
+            List<JanusGraphVertex> vertices = new ArrayList<JanusGraphVertex>(numVertices);
+            try {
+                for (int i = 0; i < numVertices; i++, count++) {
+                    final long userVertexId;
+                    switch (idStrategy) {
+                        case LOW:
+                            userVertexId = count;
+                            break;
+                        case HIGH:
+                            userVertexId = maxCount-count;
+                            break;
+                        default:
+                            throw new RuntimeException("Unsupported custom id strategy: " + idStrategy);
+                    }
+                    final long id = idAssigner.getIDManager().toVertexId(userVertexId);
+                    JanusGraphVertex next = graph.addVertex(T.id, id, "user_id", userVertexId);
+                    vertices.add(next);
+                }
+
+                //Verify that ids are set, unique and consistent with user id basis
+                for (JanusGraphVertex v : vertices) {
+                    assertTrue(v.hasId());
+                    long id = v.longId();
+                    assertTrue(id>0 && id<Long.MAX_VALUE);
+                    assertTrue(vertexIds.add(id));
+                    assertEquals((long) v.value("user_id"), idAssigner.getIDManager().fromVertexId(id));
+                }
+            } finally {
+                graph.tx().rollback();
+                graph.close();
+            }
+        }
+    }
 
 }