Merge pull request #6891 from s1ck/backport-nodesbuilder-changes

Backport NodesBuilder changes

Author: s1ck

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

@@ -42,15 +42,12 @@ public CSRGraphStoreFactory(
         super(graphProjectConfig, capabilities, loadingContext, dimensions);
     }
 
-    protected CSRGraphStore createGraphStore(
-        Nodes nodes,
-        RelationshipImportResult relationshipImportResult
-    ) {
+    protected CSRGraphStore createGraphStore(Nodes nodes, RelationshipImportResult relationshipImportResult) {
         return new GraphStoreBuilder()
             .databaseId(DatabaseId.of(loadingContext.graphDatabaseService()))
             .capabilities(capabilities)
             .schema(computeGraphSchema(nodes, relationshipImportResult))
-            .nodes(Nodes.of(nodes.idMap(), nodes.properties()))
+            .nodes(nodes)
             .relationshipImportResult(relationshipImportResult)
             .concurrency(graphProjectConfig.readConcurrency())
             .build();

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

@@ -61,8 +61,8 @@ public class ArrayIdMap extends LabeledIdMap {
 
     private final long highestNeoId;
 
-    private final HugeLongArray graphIds;
-    private final HugeSparseLongArray nodeToGraphIds;
+    private final HugeLongArray internalToOriginalIds;
+    private final HugeSparseLongArray originalToInternalIds;
 
     public static MemoryEstimation memoryEstimation() {
         return ESTIMATION;
@@ -72,26 +72,26 @@ public static MemoryEstimation memoryEstimation() {
      * initialize the map with pre-built sub arrays
      */
     public ArrayIdMap(
-        HugeLongArray graphIds,
-        HugeSparseLongArray nodeToGraphIds,
+        HugeLongArray internalToOriginalIds,
+        HugeSparseLongArray originalToInternalIds,
         LabelInformation labelInformation,
         long nodeCount,
         long highestNeoId
     ) {
         super(labelInformation, nodeCount);
-        this.graphIds = graphIds;
-        this.nodeToGraphIds = nodeToGraphIds;
+        this.internalToOriginalIds = internalToOriginalIds;
+        this.originalToInternalIds = originalToInternalIds;
         this.highestNeoId = highestNeoId;
     }
 
     @Override
     public long toMappedNodeId(long originalNodeId) {
-        return nodeToGraphIds.get(originalNodeId);
+        return originalToInternalIds.get(originalNodeId);
     }
 
     @Override
     public long toOriginalNodeId(long mappedNodeId) {
-        return graphIds.get(mappedNodeId);
+        return internalToOriginalIds.get(mappedNodeId);
     }
 
     @Override
@@ -106,7 +106,7 @@ public IdMap rootIdMap() {
 
     @Override
     public boolean contains(final long originalNodeId) {
-        return nodeToGraphIds.contains(originalNodeId);
+        return originalToInternalIds.contains(originalNodeId);
     }
 
     @Override
@@ -141,7 +141,7 @@ public Optional<FilteredIdMap> withFilteredLabels(Collection<NodeLabel> nodeLabe
 
         HugeSparseLongArray newNodeToGraphIds = ArrayIdMapBuilderOps.buildSparseIdMap(
             newNodeCount,
-            nodeToGraphIds.capacity(),
+            originalToInternalIds.capacity(),
             concurrency,
             newGraphIds
         );

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

@@ -69,8 +69,8 @@ public IdMap build(
     ) {
         adders.close();
         long nodeCount = this.size();
-        var graphIds = this.array();
-        return ArrayIdMapBuilderOps.build(graphIds, nodeCount, labelInformationBuilder, highestNodeId, concurrency);
+        HugeLongArray internalToOriginalIds = this.array();
+        return ArrayIdMapBuilderOps.build(internalToOriginalIds, nodeCount, labelInformationBuilder, highestNodeId, concurrency);
     }
 
     public HugeLongArray array() {

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

@@ -31,28 +31,28 @@
 public final class ArrayIdMapBuilderOps {
 
     static ArrayIdMap build(
-        HugeLongArray graphIds,
+        HugeLongArray internalToOriginalIds,
         long nodeCount,
         LabelInformation.Builder labelInformationBuilder,
         long highestNodeId,
         int concurrency
     ) {
         if (highestNodeId == NodesBuilder.UNKNOWN_MAX_ID) {
-            highestNodeId = graphIds.asNodeProperties().getMaxLongPropertyValue().orElse(NodesBuilder.UNKNOWN_MAX_ID);
+            highestNodeId = internalToOriginalIds.asNodeProperties().getMaxLongPropertyValue().orElse(NodesBuilder.UNKNOWN_MAX_ID);
         }
 
-        HugeSparseLongArray nodeToGraphIds = buildSparseIdMap(
+        HugeSparseLongArray originalToInternalIds = buildSparseIdMap(
             nodeCount,
             highestNodeId,
             concurrency,
-            graphIds
+            internalToOriginalIds
         );
 
-        var labelInformation = labelInformationBuilder.build(nodeCount, nodeToGraphIds::get);
+        var labelInformation = labelInformationBuilder.build(nodeCount, originalToInternalIds::get);
 
         return new ArrayIdMap(
-            graphIds,
-            nodeToGraphIds,
+            internalToOriginalIds,
+            originalToInternalIds,
             labelInformation,
             nodeCount,
             highestNodeId

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;
@@ -106,7 +104,7 @@ public static CSRGraphStore createFromGraph(
             // TODO: is it correct that we only use this for generated graphs?
             .capabilities(ImmutableStaticCapabilities.of(false))
             .schema(schema)
-            .nodes(Nodes.of(graph.idMap(), nodeProperties))
+            .nodes(ImmutableNodes.of(schema.nodeSchema(), graph.idMap(), nodeProperties))
             .relationshipImportResult(relationshipImportResult)
             .graphProperties(GraphPropertyStore.empty())
             .concurrency(concurrency)
@@ -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
         );
     }
@@ -141,25 +140,25 @@ public CSRGraphStore build() {
             ).load(ktx.internalTransaction());
 
             progressTracker.beginSubTask("Loading");
-            var idMapAndProperties = new CypherNodeLoader(
+            var nodes = new CypherNodeLoader(
                 nodeQuery(),
                 nodeCount.rows(),
                 cypherConfig,
                 loadingContext,
                 progressTracker
             ).load(ktx.internalTransaction());
 
-            var relationshipsAndProperties = new CypherRelationshipLoader(
+            var relationshipImportResult = new CypherRelationshipLoader(
                 relationshipQuery(),
-                idMapAndProperties.idMap(),
+                nodes.idMap(),
                 cypherConfig,
                 loadingContext,
                 progressTracker
             ).load(ktx.internalTransaction());
 
             var graphStore = createGraphStore(
-                idMapAndProperties,
-                relationshipsAndProperties
+                nodes,
+                relationshipImportResult
             );
 
             progressTracker.endSubTask("Loading");

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

@@ -20,19 +20,15 @@
 package org.neo4j.gds.core.loading;
 
 import org.immutables.value.Value;
-import org.neo4j.gds.PropertyMapping;
 import org.neo4j.gds.api.GraphLoaderContext;
 import org.neo4j.gds.api.PropertyState;
-import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
 import org.neo4j.gds.config.GraphProjectFromCypherConfig;
 import org.neo4j.gds.core.loading.construction.GraphFactory;
 import org.neo4j.gds.core.loading.construction.NodesBuilder;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.kernel.impl.coreapi.InternalTransaction;
 
-import java.util.Map;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 @Value.Enclosing
 class CypherNodeLoader extends CypherRecordLoader<Nodes> {
@@ -71,6 +67,7 @@ BatchLoadResult loadSingleBatch(InternalTransaction tx, int bufferSize) {
             .maxOriginalId(NodesBuilder.UNKNOWN_MAX_ID)
             .hasLabelInformation(hasLabelInformation)
             .hasProperties(!propertyColumns.isEmpty())
+            .propertyState(PropertyState.TRANSIENT)
             .build();
 
         nodeSubscriber.initialize(subscription.fieldNames(), this.nodesBuilder);
@@ -99,12 +96,7 @@ void updateCounts(BatchLoadResult result) {
 
     @Override
     Nodes result() {
-        var nodes = nodesBuilder.build(highestNodeId);
-        var idMap = nodes.idMap();
-        var nodeProperties = nodes.properties().propertyValues();
-        var nodePropertiesWithPropertyMappings = propertiesWithPropertyMappings(nodeProperties);
-
-        return Nodes.of(idMap, nodePropertiesWithPropertyMappings, PropertyState.TRANSIENT);
+        return nodesBuilder.build(highestNodeId);
     }
 
     @Override
@@ -121,16 +113,4 @@ Set<String> getReservedColumns() {
     QueryType queryType() {
         return QueryType.NODE;
     }
-
-    private static Map<PropertyMapping, NodePropertyValues> propertiesWithPropertyMappings(Map<String, NodePropertyValues> properties) {
-        return properties.entrySet()
-            .stream()
-            .collect(Collectors.toMap(
-                propertiesByKey -> PropertyMapping.of(
-                    propertiesByKey.getKey(),
-                    propertiesByKey.getValue().valueType().fallbackValue()
-                ),
-                Map.Entry::getValue
-            ));
-    }
 }

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

@@ -46,24 +46,29 @@
 import static java.util.stream.Collectors.toMap;
 import static org.neo4j.gds.core.GraphDimensions.ANY_LABEL;
 
+// TODO: should be named LoadablePropertyMappings
 final class IndexPropertyMappings {
 
-    static LoadablePropertyMappings prepareProperties(
-        GraphProjectFromStoreConfig graphProjectConfig,
-        GraphDimensions graphDimensions,
-        TransactionContext transaction
-    ) {
-        Map<NodeLabel, PropertyMappings> storeLoadedProperties = graphProjectConfig
+    static Map<NodeLabel, PropertyMappings> propertyMappings(GraphProjectFromStoreConfig graphProjectConfig) {
+        return graphProjectConfig
             .nodeProjections()
             .projections()
             .entrySet()
             .stream()
             .collect(toMap(
                 Map.Entry::getKey,
-                entry -> entry.getValue().properties()
+                entry -> entry
+                    .getValue()
+                    .properties()
             ));
+    }
 
-        return prepareLoadableProperties(graphDimensions, transaction, storeLoadedProperties);
+    static LoadablePropertyMappings prepareProperties(
+        GraphProjectFromStoreConfig graphProjectConfig,
+        GraphDimensions graphDimensions,
+        TransactionContext transaction
+    ) {
+        return prepareLoadableProperties(graphDimensions, transaction, propertyMappings(graphProjectConfig));
     }
 
     private static LoadablePropertyMappings prepareLoadableProperties(

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
         );
     }