Use the correct node count for the IdMapBuilder

Co-Authored-By: Paul Horn <paul.horn@neotechnology.com>
Co-Authored-By: Florentin Dörre <florentin.dorre@neotechnology.com>

Author: 3 people

core/src/main/java/org/neo4j/graphalgo/core/GraphDimensionsReader.java

@@ -106,13 +106,15 @@ public GraphDimensions apply(KernelTransaction transaction) throws RuntimeExcept
         Map<String, Integer> nodePropertyTokens = loadNodePropertyTokens(tokenRead);
         ResolvedPropertyMappings relProperties = loadPropertyMapping(tokenRead, setup.relationshipPropertyMappings());
 
+        // Every node will be counted once for every label it has.
         long nodeCount = labelTokenNodeLabelMappings.keyStream()
             .mapToLong(dataRead::countsForNode)
             .sum();
-        final long allNodesCount = InternalReadOps.getHighestPossibleNodeCount(dataRead, api);
+        final long highestPossibleNodeCount = InternalReadOps.getHighestPossibleNodeCount(dataRead, api);
+        // This is also an upper bound on the node count.
         long finalNodeCount = labelTokenNodeLabelMappings.keys().contains(ANY_LABEL)
-            ? allNodesCount
-            : Math.min(nodeCount, allNodesCount);
+            ? highestPossibleNodeCount
+            : Math.min(nodeCount, highestPossibleNodeCount);
 
         // TODO: this will double count relationships between distinct labels
         Map<RelationshipType, Long> relationshipCounts = relationshipProjectionMappings
@@ -133,7 +135,7 @@ public GraphDimensions apply(KernelTransaction transaction) throws RuntimeExcept
 
         return ImmutableGraphDimensions.builder()
                 .nodeCount(finalNodeCount)
-                .highestNeoId(allNodesCount)
+                .highestNeoId(highestPossibleNodeCount)
                 .maxRelCount(maxRelCount)
                 .relationshipCounts(relationshipCounts)
                 .nodeLabelIds(labelTokenNodeLabelMappings.keys())

core/src/main/java/org/neo4j/graphalgo/core/loading/IdMap.java

@@ -135,6 +135,7 @@ public IdMap withFilteredLabels(BitSet unionedBitSet, int concurrency) {
 
         SparseNodeMapping newNodeToGraphIds = IdMapBuilder.buildSparseNodeMapping(
             newGraphIds,
+            newGraphIds.size(),
             nodeToGraphIds.getCapacity(),
             concurrency,
             AllocationTracker.EMPTY

core/src/main/java/org/neo4j/graphalgo/core/loading/IdMapBuilder.java

@@ -59,21 +59,22 @@ public static IdMap build(
 
         HugeLongArray graphIds = idMapBuilder.build();
 
-        SparseNodeMapping nodeToGraphIds = buildSparseNodeMapping(graphIds, highestNodeId, concurrency, tracker);
+        SparseNodeMapping nodeToGraphIds = buildSparseNodeMapping(graphIds, idMapBuilder.size(), highestNodeId, concurrency, tracker);
         return new IdMap(graphIds, nodeToGraphIds, maybeLabelInformation, idMapBuilder.size());
     }
 
     @NotNull
     static SparseNodeMapping buildSparseNodeMapping(
         HugeLongArray graphIds,
+        long nodeCount,
         long highestNodeId,
         int concurrency,
         AllocationTracker tracker
     ) {
         SparseNodeMapping.Builder nodeMappingBuilder = SparseNodeMapping.Builder.create(highestNodeId + 1, tracker);
         ParallelUtil.readParallel(
                 concurrency,
-                graphIds.size(),
+                nodeCount,
                 Pools.DEFAULT,
                 (start, end) -> {
                     try (HugeCursor<long[]> cursor = graphIds.initCursor(graphIds.newCursor(), start, end)) {

core/src/test/java/org/neo4j/graphalgo/core/huge/NodeFilteredGraphTest.java

@@ -19,6 +19,7 @@
  */
 package org.neo4j.graphalgo.core.huge;
 
+import org.apache.commons.lang3.mutable.MutableLong;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -36,14 +37,15 @@
 import static java.util.Arrays.asList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.neo4j.graphalgo.QueryRunner.runQuery;
+import static org.neo4j.graphalgo.QueryRunner.runQueryWithRowConsumer;
 
 public class NodeFilteredGraphTest {
 
     private GraphDbApi db;
 
     private static final String DB_CYPHER =
         " CREATE" +
-        " (a:Person)," +
+        " (a:Person {name: 'Florentin'})," +
         " (b:Ignore:Person)," +
         " (c:Ignore:Person)," +
         " (d:Person)," +
@@ -62,7 +64,6 @@ void teardown() {
         db.shutdown();
     }
 
-
     @Test
     void filterDegree() {
         GraphStore graphStore = new StoreLoaderBuilder()
@@ -79,8 +80,16 @@ void filterDegree() {
             4
         );
 
-        long nodeIdOfA = graph.toMappedNodeId(0);
-        assertEquals(1L, graph.degree(nodeIdOfA));
+        String query = "MATCH(n {name: 'Florentin'}) " +
+                       "RETURN id(n) AS id";
+        MutableLong neoNodeIdForA = new MutableLong();
+
+        runQueryWithRowConsumer(db, query, row -> {
+            neoNodeIdForA.setValue(row.getNumber("id"));
+        });
+
+        long mappedNodeIdForA = graph.toMappedNodeId(neoNodeIdForA.getValue());
+        assertEquals(1L, graph.degree(mappedNodeIdForA));
     }
 
 }

core/src/test/java/org/neo4j/graphalgo/core/loading/IdMapBuilderTest.java

@@ -36,6 +36,7 @@ void buildSparseNodeMappingWithPageSizeEntries() {
 
         SparseNodeMapping hugeSparseLongArray = IdMapBuilder.buildSparseNodeMapping(
             hugeLongArray,
+            hugeLongArray.size(),
             nodeId,
             1,
             AllocationTracker.EMPTY