Move read.countsForNode/RelationshipWithoutTx to compat

DarthMax · DarthMax · commit c61a68defa9f · 2024-01-30T11:42:30.000+01:00
diff --git a/compatibility/4.4/neo4j-kernel-adapter/src/main/java/org/neo4j/gds/compat/_44/Neo4jProxyImpl.java b/compatibility/4.4/neo4j-kernel-adapter/src/main/java/org/neo4j/gds/compat/_44/Neo4jProxyImpl.java
@@ -546,6 +546,16 @@ public long getHighestPossibleRelationshipCount(
         return countByIdGenerator(idGeneratorFactory, RecordIdType.RELATIONSHIP).orElseGet(read::relationshipsGetCount);
     }
 
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
+
     @Override
     public String versionLongToString(long storeVersion) {
         return MetaDataStore.versionLongToString(storeVersion);
diff --git a/compatibility/5-common/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_5x/CommonNeo4jProxyImpl.java b/compatibility/5-common/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_5x/CommonNeo4jProxyImpl.java
@@ -177,7 +177,7 @@ public List<StoreScan<NodeLabelIndexCursor>> entityCursorScan(
         int batchSize,
         boolean allowPartitionedScan
     ) {
-        return PartitionedStoreScan.createScans(transaction, batchSize, labelIds);
+        return PartitionedStoreScan.createScans(transaction, batchSize, this,  labelIds);
     }
 
     @Override
@@ -186,7 +186,7 @@ public List<StoreScan<NodeLabelIndexCursor>> partitionedCursorScan(
         int batchSize,
         int... labelIds
     ) {
-        return PartitionedStoreScan.createScans(transaction, batchSize, labelIds);
+        return PartitionedStoreScan.createScans(transaction, batchSize, this,  labelIds);
     }
 
     @Override
@@ -264,7 +264,7 @@ public StoreScan<NodeLabelIndexCursor> nodeLabelIndexScan(
         int batchSize,
         boolean allowPartitionedScan
     ) {
-        return PartitionedStoreScan.createScans(transaction, batchSize, labelId).get(0);
+        return PartitionedStoreScan.createScans(transaction, batchSize, this, labelId).get(0);
     }
 
     @Override
diff --git a/compatibility/5-common/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_5x/PartitionedStoreScan.java b/compatibility/5-common/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_5x/PartitionedStoreScan.java
@@ -22,6 +22,7 @@
 import org.neo4j.common.EntityType;
 import org.neo4j.exceptions.KernelException;
 import org.neo4j.gds.compat.CompatExecutionContext;
+import org.neo4j.gds.compat.Neo4jProxyApi;
 import org.neo4j.gds.compat.StoreScan;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.NodeLabelIndexCursor;
@@ -49,6 +50,7 @@ public boolean reserveBatch(C cursor, CompatExecutionContext ctx) {
     public static List<StoreScan<NodeLabelIndexCursor>> createScans(
         KernelTransaction transaction,
         int batchSize,
+        Neo4jProxyApi proxy,
         int... labelIds
     ) {
         var indexDescriptor = NodeLabelIndexLookupImpl.findUsableMatchingIndex(
@@ -66,11 +68,11 @@ public static List<StoreScan<NodeLabelIndexCursor>> createScans(
         // and use that one as the driving partitioned index scan. The partitions
         // of all other partitioned index scans will be aligned to that one.
         int maxToken = labelIds[0];
-        long maxCount = read.countsForNode(labelIds[0]);
+        long maxCount = proxy.estimateNodeCount(read, labelIds[0]);
         int maxIndex = 0;
 
         for (int i = 1; i < labelIds.length; i++) {
-            long count = read.countsForNode(labelIds[i]);
+            long count = proxy.estimateNodeCount(read, labelIds[i]);
             if (count > maxCount) {
                 maxCount = count;
                 maxToken = labelIds[i];
diff --git a/compatibility/5.10/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_510/Neo4jProxyImpl.java b/compatibility/5.10/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_510/Neo4jProxyImpl.java
@@ -28,6 +28,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
 import org.neo4j.internal.kernel.api.procs.ProcedureSignature;
@@ -195,4 +196,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
             .map(DatabaseReferenceImpl.Internal::databaseId)
             .anyMatch(databaseId::equals);
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.11/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_511/Neo4jProxyImpl.java b/compatibility/5.11/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_511/Neo4jProxyImpl.java
@@ -28,6 +28,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
 import org.neo4j.internal.kernel.api.procs.ProcedureSignature;
@@ -195,4 +196,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
             .map(DatabaseReferenceImpl.Internal::databaseId)
             .anyMatch(databaseId::equals);
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.12/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_512/Neo4jProxyImpl.java b/compatibility/5.12/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_512/Neo4jProxyImpl.java
@@ -28,6 +28,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
 import org.neo4j.internal.kernel.api.procs.ProcedureSignature;
@@ -195,4 +196,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
             .map(DatabaseReferenceImpl.Internal::databaseId)
             .anyMatch(databaseId::equals);
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.13/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_513/Neo4jProxyImpl.java b/compatibility/5.13/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_513/Neo4jProxyImpl.java
@@ -30,6 +30,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.TokenSet;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
@@ -214,4 +215,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
             .map(DatabaseReferenceImpl.Internal::databaseId)
             .anyMatch(databaseId::equals);
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.14/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_514/Neo4jProxyImpl.java b/compatibility/5.14/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_514/Neo4jProxyImpl.java
@@ -30,6 +30,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.TokenSet;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
@@ -215,4 +216,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
             .map(DatabaseReferenceImpl.Internal::databaseId)
             .anyMatch(databaseId::equals);
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.15/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_515/Neo4jProxyImpl.java b/compatibility/5.15/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_515/Neo4jProxyImpl.java
@@ -30,6 +30,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.TokenSet;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
@@ -220,4 +221,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
             .map(DatabaseReferenceImpl.Internal::databaseId)
             .anyMatch(databaseId::equals);
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.estimateCountsForNode(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.estimateCountsForRelationships(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.16/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_516/Neo4jProxyImpl.java b/compatibility/5.16/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_516/Neo4jProxyImpl.java
@@ -30,6 +30,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.TokenSet;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
@@ -220,4 +221,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
             .map(DatabaseReferenceImpl.Internal::databaseId)
             .anyMatch(databaseId::equals);
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.estimateCountsForNode(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.estimateCountsForRelationships(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.6/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_56/Neo4jProxyImpl.java b/compatibility/5.6/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_56/Neo4jProxyImpl.java
@@ -32,6 +32,7 @@
 import org.neo4j.internal.kernel.api.NodeCursor;
 import org.neo4j.internal.kernel.api.NodeLabelIndexCursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.RelationshipScanCursor;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
@@ -240,4 +241,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
         var databaseManager = GraphDatabaseApiProxy.resolveDependency(databaseService, FabricDatabaseManager.class);
         return databaseManager.isFabricDatabase(GraphDatabaseApiProxy.databaseId(databaseService));
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.7/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_57/Neo4jProxyImpl.java b/compatibility/5.7/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_57/Neo4jProxyImpl.java
@@ -32,6 +32,7 @@
 import org.neo4j.internal.kernel.api.NodeCursor;
 import org.neo4j.internal.kernel.api.NodeLabelIndexCursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.RelationshipScanCursor;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
@@ -245,4 +246,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
         var databaseManager = GraphDatabaseApiProxy.resolveDependency(databaseService, FabricDatabaseManager.class);
         return databaseManager.isFabricDatabase(GraphDatabaseApiProxy.databaseId(databaseService));
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.8/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_58/Neo4jProxyImpl.java b/compatibility/5.8/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_58/Neo4jProxyImpl.java
@@ -29,6 +29,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
 import org.neo4j.internal.kernel.api.procs.ProcedureSignature;
@@ -190,4 +191,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
         var databaseManager = GraphDatabaseApiProxy.resolveDependency(databaseService, FabricDatabaseManager.class);
         return databaseManager.isFabricDatabase(GraphDatabaseApiProxy.databaseId(databaseService));
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/5.9/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_59/Neo4jProxyImpl.java b/compatibility/5.9/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_59/Neo4jProxyImpl.java
@@ -28,6 +28,7 @@
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.internal.kernel.api.Cursor;
 import org.neo4j.internal.kernel.api.PartitionedScan;
+import org.neo4j.internal.kernel.api.Read;
 import org.neo4j.internal.kernel.api.exceptions.ProcedureException;
 import org.neo4j.internal.kernel.api.procs.FieldSignature;
 import org.neo4j.internal.kernel.api.procs.ProcedureSignature;
@@ -194,4 +195,14 @@ public boolean isCompositeDatabase(GraphDatabaseService databaseService) {
             .map(DatabaseReference.Internal::databaseId)
             .anyMatch(databaseId::equals);
     }
+
+    @Override
+    public long estimateNodeCount(Read read, int label) {
+        return read.countsForNodeWithoutTxState(label);
+    }
+
+    @Override
+    public long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return read.countsForRelationshipWithoutTxState(sourceLabel, type, targetLabel);
+    }
 }
diff --git a/compatibility/api/neo4j-kernel-adapter/src/main/java/org/neo4j/gds/compat/Neo4jProxyApi.java b/compatibility/api/neo4j-kernel-adapter/src/main/java/org/neo4j/gds/compat/Neo4jProxyApi.java
@@ -225,6 +225,9 @@ ProcedureSignature procedureSignature(
 
     long getHighestPossibleRelationshipCount(Read read, @Nullable IdGeneratorFactory idGeneratorFactory);
 
+    long estimateNodeCount(Read read, int label);
+    long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type);
+
     String versionLongToString(long storeVersion);
 
     TestLog testLog();
diff --git a/compatibility/common/neo4j-kernel-adapter/src/main/java/org/neo4j/gds/compat/Neo4jProxy.java b/compatibility/common/neo4j-kernel-adapter/src/main/java/org/neo4j/gds/compat/Neo4jProxy.java
@@ -322,6 +322,14 @@ public static long getHighestPossibleRelationshipCount(Read read, @Nullable IdGe
         return IMPL.getHighestPossibleRelationshipCount(read, idGeneratorFactory);
     }
 
+    public static long estimateNodeCount(Read read, int label) {
+        return IMPL.estimateNodeCount(read, label);
+    }
+
+    public static long estimateRelationshipCount(Read read, int sourceLabel, int targetLabel, int type) {
+        return IMPL.estimateRelationshipCount(read, sourceLabel, targetLabel, type);
+    }
+
     public static String versionLongToString(long storeVersion) {
         return IMPL.versionLongToString(storeVersion);
     }
diff --git a/native-projection/src/main/java/org/neo4j/gds/projection/GraphDimensionsReader.java b/native-projection/src/main/java/org/neo4j/gds/projection/GraphDimensionsReader.java
@@ -99,7 +99,7 @@ public GraphDimensions apply(KernelTransaction transaction) throws RuntimeExcept
         Map<String, Integer> relationshipPropertyTokens = loadPropertyTokens(getRelationshipProjections().projections(), tokenRead);
 
         long nodeCount = labelTokenNodeLabelMappings.keyStream()
-            .mapToLong(dataRead::countsForNode)
+            .mapToLong(label -> Neo4jProxy.estimateNodeCount(dataRead, label))
             .sum();
         final long allNodesCount = Neo4jProxy.getHighestPossibleNodeCount(dataRead, idGeneratorFactory);
         long finalNodeCount = labelTokenNodeLabelMappings.keys().contains(ANY_LABEL)
@@ -205,8 +205,8 @@ protected Map<RelationshipType, Long> getRelationshipCountsByType(
 
     private static long relCountUpperBoundForLabelAndType(Read dataRead, int labelId, int id) {
         return Math.max(
-            dataRead.countsForRelationship(labelId, id, ANY_LABEL),
-            dataRead.countsForRelationship(ANY_LABEL, id, labelId)
+            Neo4jProxy.estimateRelationshipCount(dataRead, labelId, ANY_LABEL, id),
+            Neo4jProxy.estimateRelationshipCount(dataRead, ANY_LABEL, labelId, id)
         );
     }
 
