From 04b79e3b4d9cc9ec1ca2262aca4d299ea78e0e4d Mon Sep 17 00:00:00 2001
From: Paul Horn <paul.horn@neo4j.com>
Date: Mon, 15 Jul 2024 11:42:43 +0200
Subject: [PATCH] Support offset test graphs on block format (5.22+)

---
 .../java/org/neo4j/gds/compat/Neo4jProxy.java | 24 ++++++++----
 .../org/neo4j/gds/compat/InternalReadOps.java | 38 ++++++++-----------
 2 files changed, 31 insertions(+), 31 deletions(-)

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
index 7ac1a960f8..0f58603a74 100644
--- 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
@@ -108,7 +108,6 @@
 import java.util.stream.Stream;
 
 import static java.lang.String.format;
-import static org.neo4j.gds.compat.InternalReadOps.countByIdGenerator;
 
 public final class Neo4jProxy {
 
@@ -519,12 +518,15 @@ public static Long pageCacheMemoryValue(String value) {
     }
 
     public static long getHighestPossibleNodeCount(IdGeneratorFactory idGeneratorFactory) {
-        return countByIdGenerator(
-            idGeneratorFactory,
-            RecordIdType.NODE,
-            BlockFormat.INSTANCE.nodeType,
-            BlockFormat.INSTANCE.dynamicNodeType
-        );
+        return InternalReadOps.findValidIdGeneratorsStream(
+                idGeneratorFactory,
+                RecordIdType.NODE,
+                BlockFormat.INSTANCE.nodeType,
+                BlockFormat.INSTANCE.dynamicNodeType
+            )
+            .mapToLong(IdGenerator::getHighId)
+            .max()
+            .orElseThrow(InternalReadOps::unsupportedStoreFormatException);
     }
 
     public static long getHighestPossibleRelationshipCount(Read read) {
@@ -668,7 +670,13 @@ public static long transactionId(KernelTransaction kernelTransaction) {
     }
 
     public static void reserveNeo4jIds(IdGeneratorFactory generatorFactory, int size, CursorContext cursorContext) {
-        IdGenerator idGenerator = generatorFactory.get(RecordIdType.NODE);
+        var idGenerator = InternalReadOps.findValidIdGeneratorsStream(
+                generatorFactory,
+                RecordIdType.NODE,
+                BlockFormat.INSTANCE.nodeType,
+                BlockFormat.INSTANCE.dynamicNodeType
+            )
+            .findFirst().orElseThrow(InternalReadOps::unsupportedStoreFormatException);
 
         idGenerator.nextConsecutiveIdRange(size, false, cursorContext);
     }
diff --git a/neo4j-adapter/src/main/java/org/neo4j/gds/compat/InternalReadOps.java b/neo4j-adapter/src/main/java/org/neo4j/gds/compat/InternalReadOps.java
index c08e296069..c7903313e7 100644
--- a/neo4j-adapter/src/main/java/org/neo4j/gds/compat/InternalReadOps.java
+++ b/neo4j-adapter/src/main/java/org/neo4j/gds/compat/InternalReadOps.java
@@ -24,41 +24,33 @@
 import org.neo4j.internal.id.IdGeneratorFactory;
 import org.neo4j.internal.id.IdType;
 
-import java.util.OptionalLong;
+import java.util.Arrays;
+import java.util.stream.Stream;
 
 public final class InternalReadOps {
 
-    public static long countByIdGenerator(
+    public static Stream<IdGenerator> findValidIdGeneratorsStream(
         @Nullable IdGeneratorFactory idGeneratorFactory,
         IdType... idTypes
     ) {
-        long highestId = Long.MIN_VALUE;
-        for (IdType idType : idTypes) {
-            final OptionalLong count = countByIdGenerator(idGeneratorFactory, idType);
-            if (count.isPresent()) {
-                highestId = Math.max(highestId, count.getAsLong());
-            }
-        }
-        if (highestId == Long.MIN_VALUE) {
-            throw new IllegalStateException(
-                "Unsupported store format for GDS; GDS cannot read data from this database. " +
-                    "Please try to use Cypher projection instead.");
+        if (idGeneratorFactory == null || idTypes.length == 0) {
+            return Stream.empty();
         }
-        return highestId;
-    }
-
-    private static OptionalLong countByIdGenerator(@Nullable IdGeneratorFactory idGeneratorFactory, @Nullable IdType idType) {
-        if (idGeneratorFactory != null && idType != null) {
+        return Arrays.stream(idTypes).mapMulti((idType, downstream) -> {
             try {
-                final IdGenerator idGenerator = idGeneratorFactory.get(idType);
+                var idGenerator = idGeneratorFactory.get(idType);
                 if (idGenerator != null) {
-                    // getHighId returns the highestId + 1, which is actually a count
-                    return OptionalLong.of(idGenerator.getHighId());
+                    downstream.accept(idGenerator);
                 }
             } catch (Exception ignored) {
             }
-        }
-        return OptionalLong.empty();
+        });
+    }
+
+    public static IllegalStateException unsupportedStoreFormatException() {
+        return new IllegalStateException(
+            "Unsupported store format for GDS; GDS cannot read data from this database. " +
+                "Please try to use Cypher projection instead.");
     }
 
     private InternalReadOps() {