Merge pull request #11732 from soerenreichardt/arrow-graphstore-export

Arrow graphstore export

Author: soerenreichardt

core-api/src/main/java/org/neo4j/gds/api/GraphStoreWithTopology.java

@@ -0,0 +1,26 @@
+/*
+ * 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.api;
+
+import org.neo4j.gds.RelationshipType;
+
+public interface GraphStoreWithTopology extends GraphStore {
+    Topology getTopology(RelationshipType relationshipType);
+}

core/src/main/java/org/neo4j/gds/api/GraphStoreAdapter.java

@@ -38,7 +38,7 @@
 import java.util.Optional;
 import java.util.Set;
 
-public abstract class GraphStoreAdapter implements GraphStore {
+public abstract class GraphStoreAdapter implements GraphStoreWithTopology {
 
     private final GraphStore graphStore;
 
@@ -318,4 +318,14 @@ public CompositeRelationshipIterator getCompositeRelationshipIterator(
     ) {
         return graphStore.getCompositeRelationshipIterator(relationshipType, propertyKeys);
     }
+
+    @Override
+    public Topology getTopology(RelationshipType relationshipType) {
+        if (graphStore instanceof GraphStoreWithTopology gWT) {
+            return gWT.getTopology(relationshipType);
+        }
+        throw new UnsupportedOperationException(
+            "GraphStore does not support getting topology for relationship type: " + relationshipType
+        );
+    }
 }

core/src/main/java/org/neo4j/gds/core/loading/CSRGraphStore.java

@@ -30,7 +30,7 @@
 import org.neo4j.gds.api.DatabaseInfo;
 import org.neo4j.gds.api.FilteredIdMap;
 import org.neo4j.gds.api.GraphCharacteristics;
-import org.neo4j.gds.api.GraphStore;
+import org.neo4j.gds.api.GraphStoreWithTopology;
 import org.neo4j.gds.api.IdMap;
 import org.neo4j.gds.api.PropertyState;
 import org.neo4j.gds.api.Topology;
@@ -77,7 +77,7 @@
 import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
 
 @Value.Style(typeBuilder = "GraphStoreBuilder")
-public final class CSRGraphStore implements GraphStore {
+public final class CSRGraphStore implements GraphStoreWithTopology {
 
     private final Concurrency concurrency;
 
@@ -474,6 +474,7 @@ public DeletionResult deleteRelationships(RelationshipType relationshipType) {
         }));
     }
 
+    @Override
     public Topology getTopology(RelationshipType relationshipType) {
         return relationships.get(relationshipType).topology();
     }