Merge pull request #5184 from soerenreichardt/arrow-export-array-properties

Export array node properties via Flight

Author: knutwalker

algo/src/main/java/org/neo4j/gds/louvain/Louvain.java

@@ -25,7 +25,6 @@
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.IdMap;
 import org.neo4j.gds.api.RelationshipIterator;
-import org.neo4j.gds.api.properties.nodes.LongNodePropertyValues;
 import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
 import org.neo4j.gds.beta.modularity.ImmutableModularityOptimizationStreamConfig;
 import org.neo4j.gds.beta.modularity.ModularityOptimization;
@@ -35,12 +34,11 @@
 import org.neo4j.gds.core.concurrency.ParallelUtil;
 import org.neo4j.gds.core.loading.construction.GraphFactory;
 import org.neo4j.gds.core.loading.construction.RelationshipsBuilder;
+import org.neo4j.gds.core.utils.OriginalIdNodePropertyValues;
 import org.neo4j.gds.core.utils.paged.HugeLongArray;
 import org.neo4j.gds.core.utils.partition.Partition;
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
-import org.neo4j.values.storable.Value;
-import org.neo4j.values.storable.Values;
 
 import java.util.Optional;
 import java.util.OptionalLong;
@@ -98,7 +96,13 @@ public Louvain compute() {
             long maxCommunityId = buildDendrogram(workingGraph, ranLevels, modularityOptimization);
 
             workingGraph = summarizeGraph(workingGraph, modularityOptimization, maxCommunityId);
-            nextSeedingValues = new OriginalIdNodePropertyValues(workingGraph);
+            nextSeedingValues = new OriginalIdNodePropertyValues(workingGraph) {
+                @Override
+                public OptionalLong getMaxLongPropertyValue() {
+                    // We want to use the maxSeedCommunity with value 0 in all subsequent iterations
+                    return OptionalLong.empty();
+                }
+            };
 
             if (workingGraph.nodeCount() == oldNodeCount
                 || workingGraph.nodeCount() == 1
@@ -269,34 +273,6 @@ public void release() {
         this.rootGraph.releaseTopology();
     }
 
-    static class OriginalIdNodePropertyValues implements LongNodePropertyValues {
-        private final Graph graph;
-
-        public OriginalIdNodePropertyValues(Graph graph) {
-            this.graph = graph;
-        }
-
-        @Override
-        public long longValue(long nodeId) {
-            return graph.toOriginalNodeId(nodeId);
-        }
-
-        @Override
-        public Value value(long nodeId) {
-            return Values.longValue(longValue(nodeId));
-        }
-
-        @Override
-        public OptionalLong getMaxLongPropertyValue() {
-            return OptionalLong.empty();
-        }
-
-        @Override
-        public long size() {
-            return graph.nodeCount();
-        }
-    }
-
     static final class RelationshipCreator implements Runnable {
 
         private final RelationshipsBuilder relationshipsBuilder;

…o4j/gds/core/IdentityPropertyValues.java …s/core/utils/IdentityPropertyValues.javatest-utils/src/main/java/org/neo4j/gds/core/IdentityPropertyValues.java renamed to core/src/main/java/org/neo4j/gds/core/utils/IdentityPropertyValues.java test-utils/src/main/java/org/neo4j/gds/core/IdentityPropertyValues.java renamed to core/src/main/java/org/neo4j/gds/core/utils/IdentityPropertyValues.java

@@ -17,7 +17,7 @@
  * You should have received a copy of the GNU General Public License
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
-package org.neo4j.gds.core;
+package org.neo4j.gds.core.utils;
 
 import org.neo4j.gds.api.properties.nodes.LongNodePropertyValues;
 

core/src/main/java/org/neo4j/gds/core/utils/OriginalIdNodePropertyValues.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.core.utils;
+
+import org.neo4j.gds.api.IdMap;
+import org.neo4j.gds.api.properties.nodes.LongNodePropertyValues;
+
+import java.util.OptionalLong;
+
+public class OriginalIdNodePropertyValues implements LongNodePropertyValues {
+    private final IdMap idMap;
+
+    public OriginalIdNodePropertyValues(IdMap idMap) {
+        this.idMap = idMap;
+    }
+
+    @Override
+    public long longValue(long nodeId) {
+        return idMap.toOriginalNodeId(nodeId);
+    }
+
+    @Override
+    public OptionalLong getMaxLongPropertyValue() {
+        return OptionalLong.of(idMap.highestNeoId());
+    }
+
+    @Override
+    public long size() {
+        return idMap.nodeCount();
+    }
+}

core/src/test/java/org/neo4j/gds/core/loading/GraphStoreCatalogTest.java

@@ -43,7 +43,7 @@
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
-import static org.neo4j.gds.extension.GdlSupportExtension.DATABASE_ID;
+import static org.neo4j.gds.extension.GdlSupportPerMethodExtension.DATABASE_ID;
 
 @GdlExtension
 class GraphStoreCatalogTest {

proc/catalog/src/test/java/org/neo4j/gds/catalog/GraphStreamNodePropertiesProcTest.java

@@ -32,8 +32,8 @@
 import org.neo4j.gds.PropertyMappings;
 import org.neo4j.gds.api.GraphStore;
 import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
-import org.neo4j.gds.core.IdentityPropertyValues;
 import org.neo4j.gds.core.loading.GraphStoreCatalog;
+import org.neo4j.gds.core.utils.IdentityPropertyValues;
 import org.neo4j.gds.functions.AsNodeFunc;
 
 import java.util.Map;

proc/catalog/src/test/java/org/neo4j/gds/catalog/GraphWriteNodePropertiesProcTest.java

@@ -35,8 +35,8 @@
 import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
 import org.neo4j.gds.compat.Neo4jProxy;
 import org.neo4j.gds.compat.TestLog;
-import org.neo4j.gds.core.IdentityPropertyValues;
 import org.neo4j.gds.core.loading.GraphStoreCatalog;
+import org.neo4j.gds.core.utils.IdentityPropertyValues;
 import org.neo4j.gds.core.write.NativeNodePropertiesExporterBuilder;
 import org.neo4j.gds.degree.DegreeCentralityMutateProc;
 import org.neo4j.gds.pagerank.PageRankMutateProc;

test-utils/src/main/java/org/neo4j/gds/TestSupport.java

@@ -37,7 +37,7 @@
 import org.neo4j.gds.core.GraphDimensions;
 import org.neo4j.gds.core.loading.construction.GraphFactory;
 import org.neo4j.gds.core.utils.mem.MemoryEstimation;
-import org.neo4j.gds.extension.GdlSupportExtension;
+import org.neo4j.gds.extension.GdlSupportPerMethodExtension;
 import org.neo4j.gds.extension.IdFunction;
 import org.neo4j.gds.extension.TestGraph;
 import org.neo4j.gds.gdl.GdlFactory;
@@ -155,7 +155,7 @@ public static TestGraph gdl(
             .builder()
             .nodeIdFunction(idSupplier.orElse(new OffsetIdSupplier(0L)))
             .graphProjectConfig(config)
-            .namedDatabaseId(GdlSupportExtension.DATABASE_ID)
+            .namedDatabaseId(GdlSupportPerMethodExtension.DATABASE_ID)
             .build();
 
         return new TestGraph(gdlFactory.build().getUnion(), gdlFactory::nodeId, graphName);

…j/gds/extension/GdlSupportExtension.java …s/extension/BaseGdlSupportExtension.javatest-utils/src/main/java/org/neo4j/gds/extension/GdlSupportExtension.java renamed to test-utils/src/main/java/org/neo4j/gds/extension/BaseGdlSupportExtension.java test-utils/src/main/java/org/neo4j/gds/extension/GdlSupportExtension.java renamed to test-utils/src/main/java/org/neo4j/gds/extension/BaseGdlSupportExtension.java

@@ -20,8 +20,6 @@
 package org.neo4j.gds.extension;
 
 import org.immutables.value.Value;
-import org.junit.jupiter.api.extension.AfterEachCallback;
-import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionConfigurationException;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.neo4j.gds.Orientation;
@@ -52,18 +50,16 @@
 import static org.neo4j.gds.extension.ExtensionUtil.getStringValueOfField;
 import static org.neo4j.gds.extension.ExtensionUtil.setField;
 
-public class GdlSupportExtension implements BeforeEachCallback, AfterEachCallback {
+abstract class BaseGdlSupportExtension {
 
     public static final NamedDatabaseId DATABASE_ID = DatabaseIdFactory.from("GDL", UUID.fromString("42-42-42-42-42"));
 
-    @Override
-    public void beforeEach(ExtensionContext context) {
+    void beforeAction(ExtensionContext context) {
         Class<?> requiredTestClass = context.getRequiredTestClass();
         gdlGraphs(requiredTestClass).forEach(setup -> injectGraphStore(setup, context));
     }
 
-    @Override
-    public void afterEach(ExtensionContext context) {
+    void afterAction(ExtensionContext context) {
         GraphStoreCatalog.removeAllLoadedGraphs();
     }
 
@@ -73,7 +69,7 @@ private static Collection<GdlGraphSetup> gdlGraphs(Class<?> testClass) {
         do {
             stream(testClass.getDeclaredFields())
                 .filter(f -> f.isAnnotationPresent(GdlGraph.class) || f.isAnnotationPresent(GdlGraphs.class))
-                .flatMap(GdlSupportExtension::gdlGraphsForField)
+                .flatMap(BaseGdlSupportExtension::gdlGraphsForField)
                 .peek(setup -> {
                     if (setups.contains(setup)) {
                         throw new ExtensionConfigurationException(String.format(
@@ -158,7 +154,13 @@ private static void injectGraphStore(GdlGraphSetup gdlGraphSetup, ExtensionConte
         });
     }
 
-    private static <T> void injectInstance(Object testInstance, String graphNamePrefix, T instance, Class<T> clazz, String suffix) {
+    private static <T> void injectInstance(
+        Object testInstance,
+        String graphNamePrefix,
+        T instance,
+        Class<T> clazz,
+        String suffix
+    ) {
         Stream.<Class<?>>iterate(testInstance.getClass(), Objects::nonNull, Class::getSuperclass)
             .flatMap(c -> Arrays.stream(c.getDeclaredFields()))
             .filter(field -> field.getType() == clazz)

test-utils/src/main/java/org/neo4j/gds/extension/GdlExtension.java

@@ -30,6 +30,6 @@
 @Inherited
 @Target({ElementType.TYPE, ElementType.METHOD})
 @Retention(RetentionPolicy.RUNTIME)
-@ExtendWith(GdlSupportExtension.class)
+@ExtendWith(GdlSupportPerMethodExtension.class)
 public @interface GdlExtension {
 }

test-utils/src/main/java/org/neo4j/gds/extension/GdlPerClassExtension.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.extension;
+
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Inherited
+@Target({ElementType.TYPE, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+@ExtendWith(GdlSupportPerClassExtension.class)
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public @interface GdlPerClassExtension {
+}