Merge pull request #8855 from knutwalker/2.6-native-block

Try record and block formats when trying to find the highest node count

Author: knutwalker

compatibility/4.4/neo4j-kernel-adapter/src/main/java/org/neo4j/gds/compat/_44/Neo4jProxyImpl.java

@@ -537,14 +537,14 @@ public ProcedureSignature procedureSignature(
     public long getHighestPossibleNodeCount(
         Read read, IdGeneratorFactory idGeneratorFactory
     ) {
-        return countByIdGenerator(idGeneratorFactory, RecordIdType.NODE).orElseGet(read::nodesGetCount);
+        return countByIdGenerator(idGeneratorFactory, RecordIdType.NODE, RecordIdType.NODE);
     }
 
     @Override
     public long getHighestPossibleRelationshipCount(
         Read read, IdGeneratorFactory idGeneratorFactory
     ) {
-        return countByIdGenerator(idGeneratorFactory, RecordIdType.RELATIONSHIP).orElseGet(read::relationshipsGetCount);
+        return countByIdGenerator(idGeneratorFactory, RecordIdType.RELATIONSHIP, RecordIdType.RELATIONSHIP);
     }
 
     @Override

compatibility/5-common/neo4j-kernel-adapter/src/main/java17/org/neo4j/gds/compat/_5x/CommonNeo4jProxyImpl.java

@@ -124,6 +124,7 @@
 import java.nio.file.Path;
 import java.util.List;
 import java.util.Locale;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
@@ -452,14 +453,45 @@ public Long pageCacheMemoryValue(String value) {
     public long getHighestPossibleNodeCount(
         Read read, IdGeneratorFactory idGeneratorFactory
     ) {
-        return countByIdGenerator(idGeneratorFactory, RecordIdType.NODE).orElseGet(read::nodesGetCount);
+        return countByIdGenerator(idGeneratorFactory, RecordIdType.NODE, BlockFormat.INSTANCE.nodeType);
     }
 
     @Override
     public long getHighestPossibleRelationshipCount(
         Read read, IdGeneratorFactory idGeneratorFactory
     ) {
-        return countByIdGenerator(idGeneratorFactory, RecordIdType.RELATIONSHIP).orElseGet(read::relationshipsGetCount);
+        return countByIdGenerator(idGeneratorFactory, RecordIdType.RELATIONSHIP, BlockFormat.INSTANCE.relationshipType);
+    }
+
+    private static final class BlockFormat {
+        private static final BlockFormat INSTANCE = new BlockFormat();
+
+        private final org.neo4j.internal.id.IdType nodeType;
+        private final org.neo4j.internal.id.IdType relationshipType;
+
+        BlockFormat() {
+            org.neo4j.internal.id.IdType nodeType = null;
+            org.neo4j.internal.id.IdType relationshipType = null;
+
+            try {
+                var blockIdType = Class.forName("com.neo4j.internal.blockformat.BlockIdType");
+                var blockTypes = Objects.requireNonNull(blockIdType.getEnumConstants());
+                for (Object blockType : blockTypes) {
+                    var type = (Enum<?>) blockType;
+                    if (type.name().equals("NODE")) {
+                        nodeType = (org.neo4j.internal.id.IdType) type;
+                    } else if (type.name().equals("RELATIONSHIP")) {
+                        relationshipType = (org.neo4j.internal.id.IdType) type;
+                    }
+                }
+            } catch (ClassNotFoundException | NullPointerException | ClassCastException e) {
+                nodeType = null;
+                relationshipType = null;
+            }
+
+            this.nodeType = Objects.requireNonNullElse(nodeType, RecordIdType.NODE);
+            this.relationshipType = Objects.requireNonNullElse(relationshipType, RecordIdType.RELATIONSHIP);
+        }
     }
 
     @Override

native-projection/src/main/java/org/neo4j/gds/projection/BufferedNodeConsumer.java

@@ -32,7 +32,7 @@
 import java.util.Optional;
 
 import static org.neo4j.gds.utils.GdsFeatureToggles.SKIP_ORPHANS;
-import static org.neo4j.kernel.impl.store.record.AbstractBaseRecord.NO_ID;
+import static org.neo4j.internal.kernel.api.Read.NO_ID;
 
 public final class BufferedNodeConsumer implements StoreScanner.RecordConsumer<NodeReference> {
 

neo4j-adapter/src/main/java/org/neo4j/gds/compat/InternalReadOps.java

@@ -28,7 +28,19 @@
 
 public final class InternalReadOps {
 
-    public static OptionalLong countByIdGenerator(@Nullable IdGeneratorFactory idGeneratorFactory, IdType idType) {
+    public static long countByIdGenerator(
+        @Nullable IdGeneratorFactory idGeneratorFactory,
+        IdType idType,
+        IdType idType2
+    ) {
+        return countByIdGenerator(idGeneratorFactory, idType)
+            .orElseGet(() -> countByIdGenerator(idGeneratorFactory, idType2)
+                .orElseThrow(() -> new IllegalStateException(
+                    "Unsupported store format for GDS; GDS cannot read data from this database. " +
+                        "Please try to use Cypher projection instead.")));
+    }
+
+    private static OptionalLong countByIdGenerator(@Nullable IdGeneratorFactory idGeneratorFactory, IdType idType) {
         if (idGeneratorFactory != null) {
             try {
                 final IdGenerator idGenerator = idGeneratorFactory.get(idType);