Use Nodes#nodeSchema in arrow import

Co-authored-by: Martin Junghanns <martin.junghanns@neotechnology.com>

Author: soerenreichardt

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

@@ -20,7 +20,6 @@
 package org.neo4j.gds.core.loading;
 
 import org.jetbrains.annotations.NotNull;
-import org.neo4j.gds.NodeLabel;
 import org.neo4j.gds.RelationshipType;
 import org.neo4j.gds.api.DatabaseId;
 import org.neo4j.gds.api.Graph;
@@ -43,7 +42,6 @@
 import org.neo4j.gds.core.huge.HugeGraph;
 import org.neo4j.values.storable.NumberType;
 
-import java.util.Collection;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
@@ -203,25 +201,12 @@ public static void extractNodeProperties(
         nodeImportResultBuilder.properties(propertyStoreBuilder.build());
     }
 
+    // TODO: remove this method
     public static GraphSchema computeGraphSchema(
         Nodes nodes,
-        Function<NodeLabel, Collection<String>> propertiesByLabel,
         RelationshipImportResult relationshipImportResult
     ) {
-        var nodeProperties = nodes.properties().properties();
-
-        var nodeSchema = NodeSchema.empty();
-        for (var label : nodes.idMap().availableNodeLabels()) {
-            var entry = nodeSchema.getOrCreateLabel(label);
-            for (var propertyKey : propertiesByLabel.apply(label)) {
-                entry.addProperty(
-                    propertyKey,
-                    nodeProperties.get(propertyKey).propertySchema()
-                );
-            }
-        }
-        nodes.idMap().availableNodeLabels().forEach(nodeSchema::getOrCreateLabel);
-
+        var nodeSchema = nodes.schema();
         var relationshipSchema = RelationshipSchema.empty();
 
         relationshipImportResult.importResults().forEach(((relationshipType, singleTypeRelationshipImportResult) -> {

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

@@ -124,7 +124,6 @@ protected GraphSchema computeGraphSchema(
     ) {
         return CSRGraphStoreUtil.computeGraphSchema(
             nodes,
-            (__) -> nodes.properties().keySet(),
             relationshipImportResult
         );
     }

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

@@ -23,6 +23,7 @@
 import org.neo4j.gds.api.IdMap;
 import org.neo4j.gds.api.PartialIdMap;
 import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
+import org.neo4j.gds.api.schema.NodeSchema;
 import org.neo4j.gds.core.loading.construction.GraphFactory;
 import org.neo4j.gds.core.loading.construction.NodeLabelToken;
 import org.neo4j.gds.core.loading.construction.NodesBuilder;
@@ -110,6 +111,8 @@ public interface HighLimitIdMapAndProperties {
 
         PartialIdMap intermediateIdMap();
 
+        NodeSchema schema();
+
         Optional<Map<String, NodePropertyValues>> nodeProperties();
     }
 
@@ -141,6 +144,7 @@ public OptionalLong rootNodeCount() {
             .builder()
             .idMap(idMap)
             .intermediateIdMap(partialIdMap)
+            .schema(nodes.schema())
             .nodeProperties(nodes.properties().propertyValues())
             .build();
     }

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

@@ -301,7 +301,6 @@ protected GraphSchema computeGraphSchema(
     ) {
         return CSRGraphStoreUtil.computeGraphSchema(
             nodes,
-            (label) -> storeConfig.nodeProjections().projections().get(label).properties().propertyKeys(),
             relationshipImportResult
         );
     }

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

@@ -24,6 +24,7 @@
 import org.neo4j.gds.annotation.ValueClass;
 import org.neo4j.gds.api.IdMap;
 import org.neo4j.gds.api.PropertyState;
+import org.neo4j.gds.api.properties.nodes.ImmutableNodeProperty;
 import org.neo4j.gds.api.properties.nodes.NodeProperty;
 import org.neo4j.gds.api.properties.nodes.NodePropertyStore;
 import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
@@ -54,6 +55,23 @@ static Nodes of(IdMap idmap, NodePropertyStore nodePropertyStore) {
         return ImmutableNodes.of(NodeSchema.empty(), idmap, nodePropertyStore);
     }
 
+    static Nodes of(IdMap idMap, NodeSchema nodeSchema, Map<String, NodePropertyValues> properties, PropertyState propertyState) {
+        var nodePropertyStoreBuilder = NodePropertyStore.builder();
+
+        nodeSchema.availableLabels().forEach(nodeLabel -> {
+            var propertiesForLabel = nodeSchema.get(nodeLabel).properties();
+            propertiesForLabel.forEach((propertyKey, propertySchema) -> nodePropertyStoreBuilder.putProperty(propertyKey,
+                ImmutableNodeProperty
+                    .builder()
+                    .propertySchema(propertySchema)
+                    .values(properties.get(propertyKey))
+                    .build()
+            ));
+        });
+
+        return ImmutableNodes.of(nodeSchema, idMap, nodePropertyStoreBuilder.build());
+    }
+
     static Nodes of(IdMap idMap, Map<PropertyMapping, NodePropertyValues> properties, PropertyState propertyState) {
         NodePropertyStore.Builder builder = NodePropertyStore.builder();
         properties.forEach((mapping, nodeProperties) -> builder.putProperty(