Make property schema building thread local

Co-authored-by: Sören Reichardt <soren.reichardt@neotechnology.com>

Author: s1ck

core/src/main/java/org/neo4j/gds/core/loading/construction/GraphFactory.java

@@ -121,7 +121,7 @@ static NodesBuilder nodesBuilder(
             maxOriginalId,
             threadCount,
             TokenToNodeLabel.lazy(),
-            NodeLabelTokenToPropertyKeys.lazy(),
+            NodeLabelTokenToPropertyKeys::lazy,
             new ConcurrentHashMap<>(),
             idMapBuilder,
             labelInformation,
@@ -162,7 +162,7 @@ private static NodesBuilder fromSchema(
             maxOriginalId,
             concurrency,
             TokenToNodeLabel.fixed(elementIdentifierLabelTokenMapping, labelTokenNodeLabelMapping),
-            NodeLabelTokenToPropertyKeys.fixed(nodeSchema),
+            () -> NodeLabelTokenToPropertyKeys.fixed(nodeSchema),
             new ConcurrentHashMap<>(propertyBuildersByPropertyKey),
             idMapBuilder,
             hasLabelInformation,

core/src/main/java/org/neo4j/gds/core/loading/construction/NodeLabelTokenToPropertyKeys.java

@@ -24,10 +24,10 @@
 import org.neo4j.gds.api.schema.PropertySchema;
 import org.neo4j.gds.utils.StringJoining;
 
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
 abstract class NodeLabelTokenToPropertyKeys {
@@ -70,6 +70,25 @@ abstract Map<String, PropertySchema> propertySchemas(
         Map<String, PropertySchema> importPropertySchemas
     );
 
+    static NodeLabelTokenToPropertyKeys merge(
+        NodeLabelTokenToPropertyKeys left,
+        NodeLabelTokenToPropertyKeys right,
+        Map<String, PropertySchema> importPropertySchemas
+    ) {
+        var merge = NodeLabelTokenToPropertyKeys.lazy();
+
+        left.nodeLabels().forEach(nodeLabel -> {
+            var propertyKeys = left.propertySchemas(nodeLabel, importPropertySchemas).keySet();
+            merge.add(NodeLabelTokens.ofNodeLabel(nodeLabel), propertyKeys);
+        });
+        right.nodeLabels().forEach(nodeLabel -> {
+            var propertyKeys = right.propertySchemas(nodeLabel, importPropertySchemas).keySet();
+            merge.add(NodeLabelTokens.ofNodeLabel(nodeLabel), propertyKeys);
+        });
+
+        return merge;
+    }
+
     private static class Fixed extends NodeLabelTokenToPropertyKeys {
 
         private final NodeSchema nodeSchema;
@@ -136,10 +155,10 @@ Map<String, PropertySchema> propertySchemas(
 
     private static class Lazy extends NodeLabelTokenToPropertyKeys {
 
-        private final ConcurrentHashMap<NodeLabelToken, Set<String>> labelToPropertyKeys;
+        private final Map<NodeLabelToken, Set<String>> labelToPropertyKeys;
 
         Lazy() {
-            this.labelToPropertyKeys = new ConcurrentHashMap<>();
+            this.labelToPropertyKeys = new HashMap<>();
         }
 
         @Override

core/src/main/java/org/neo4j/gds/core/loading/construction/NodesBuilder.java

@@ -50,6 +50,7 @@
 import org.neo4j.values.virtual.MapValue;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -58,6 +59,7 @@
 import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Function;
 import java.util.function.LongPredicate;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import static java.util.stream.Collectors.toMap;
@@ -81,13 +83,11 @@ public final class NodesBuilder {
 
     private final ConcurrentMap<String, NodePropertiesFromStoreBuilder> propertyBuildersByPropertyKey;
 
-    private final NodeLabelTokenToPropertyKeys nodeLabelTokenToPropertyKeys;
-
     NodesBuilder(
         long maxOriginalId,
         int concurrency,
         TokenToNodeLabel tokenToNodeLabel,
-        NodeLabelTokenToPropertyKeys nodeLabelTokenToPropertyKeys,
+        Supplier<NodeLabelTokenToPropertyKeys> nodeLabelTokenToPropertyKeysSupplier,
         ConcurrentMap<String, NodePropertiesFromStoreBuilder> propertyBuildersByPropertyKey,
         IdMapBuilder idMapBuilder,
         boolean hasLabelInformation,
@@ -97,7 +97,6 @@ public final class NodesBuilder {
     ) {
         this.maxOriginalId = maxOriginalId;
         this.concurrency = concurrency;
-        this.nodeLabelTokenToPropertyKeys = nodeLabelTokenToPropertyKeys;
         this.idMapBuilder = idMapBuilder;
         this.propertyStates = propertyStates;
         this.labelInformationBuilder = !hasLabelInformation
@@ -128,7 +127,7 @@ public final class NodesBuilder {
                 hasLabelInformation,
                 hasProperties,
                 tokenToNodeLabel,
-                nodeLabelTokenToPropertyKeys,
+                nodeLabelTokenToPropertyKeysSupplier.get(),
                 propertyBuilderFn
             )
         );
@@ -202,12 +201,15 @@ public Nodes build() {
     public Nodes build(long highestNeoId) {
         // Flush remaining buffer contents
         this.threadLocalBuilder.forEach(ThreadLocalBuilder::flush);
+        // Collect token to property keys for final merge
+        var nodeLabelTokenToPropertyKeysList = new ArrayList<NodeLabelTokenToPropertyKeys>();
+        this.threadLocalBuilder.forEach(tlb -> nodeLabelTokenToPropertyKeysList.add(tlb.nodeLabelTokenToPropertyKeys));
         // Clean up resources held by local builders
         this.threadLocalBuilder.close();
 
         var idMap = this.idMapBuilder.build(labelInformationBuilder, highestNeoId, concurrency);
         var nodeProperties = buildProperties(idMap);
-        var nodeSchema = buildNodeSchema(idMap, nodeProperties);
+        var nodeSchema = buildNodeSchema(idMap, nodeLabelTokenToPropertyKeysList, nodeProperties);
         var nodePropertyStore = NodePropertyStore.builder().properties(nodeProperties).build();
 
         return ImmutableNodes.builder()
@@ -217,7 +219,11 @@ public Nodes build(long highestNeoId) {
             .build();
     }
 
-    private NodeSchema buildNodeSchema(IdMap idMap, Map<String, NodeProperty> nodeProperties) {
+    private NodeSchema buildNodeSchema(
+        IdMap idMap,
+        Collection<NodeLabelTokenToPropertyKeys> localNodeLabelTokenToPropertyKeys,
+        Map<String, NodeProperty> nodeProperties
+    ) {
         var nodeSchema = NodeSchema.empty();
         var importPropertySchemas = nodeProperties
             .entrySet()
@@ -227,12 +233,19 @@ private NodeSchema buildNodeSchema(IdMap idMap, Map<String, NodeProperty> nodePr
         // consider node labels without properties
         var nodeLabels = new HashSet<>(idMap.availableNodeLabels());
         // and also node labels with associated properties
-        nodeLabels.addAll(nodeLabelTokenToPropertyKeys.nodeLabels());
-
+        localNodeLabelTokenToPropertyKeys.forEach(mapping -> nodeLabels.addAll(mapping.nodeLabels()));
+        // merge into a global mapping
+        var globalNodeLabelTokenToPropertyKeys = localNodeLabelTokenToPropertyKeys
+            .stream()
+            .reduce(
+                NodeLabelTokenToPropertyKeys.lazy(),
+                (left, right) -> NodeLabelTokenToPropertyKeys.merge(left, right, importPropertySchemas)
+            );
+        // and construct final node property schema
         nodeLabels.forEach(nodeLabel -> {
             nodeSchema.addLabel(
                 nodeLabel,
-                this.nodeLabelTokenToPropertyKeys.propertySchemas(nodeLabel, importPropertySchemas)
+                globalNodeLabelTokenToPropertyKeys.propertySchemas(nodeLabel, importPropertySchemas)
             );
         });
 
@@ -331,6 +344,10 @@ private static class ThreadLocalBuilder implements AutoCloseable {
             this.batchNodeProperties = new ArrayList<>(buffer.capacity());
         }
 
+        NodeLabelTokenToPropertyKeys nodeLabelTokenToPropertyKeys() {
+            return this.nodeLabelTokenToPropertyKeys;
+        }
+
         public void addNode(long originalId, NodeLabelToken nodeLabels) {
             if (!seenNodeIdPredicate.test(originalId)) {
                 long[] labels = getOrCreateLabelTokens(nodeLabels);

core/src/test/java/org/neo4j/gds/core/loading/construction/NodeLabelTokenToPropertyKeysTest.java

@@ -199,5 +199,29 @@ void testNodeLabelsFixed() {
 
         assertThat(mapping.nodeLabels()).containsExactlyInAnyOrder(NodeLabel.ALL_NODES, NodeLabel.of("A"), NodeLabel.of("B"), NodeLabel.of("C"));
     }
+
+    @Test
+    void testMerge() {
+        var importPropertySchemas = Map.of(
+            "foo", PropertySchema.of("foo", ValueType.LONG, DefaultValue.forLong(), PropertyState.TRANSIENT),
+            "bar", PropertySchema.of("bar", ValueType.DOUBLE, DefaultValue.forDouble(), PropertyState.PERSISTENT)
+        );
+
+        var left = NodeLabelTokenToPropertyKeys.lazy();
+        left.add(NodeLabelTokens.ofStrings("A"), List.of("foo"));
+
+        var right = NodeLabelTokenToPropertyKeys.lazy();
+        right.add(NodeLabelTokens.ofStrings("B"), List.of("bar"));
+
+        var mapping = NodeLabelTokenToPropertyKeys.merge(left, right, importPropertySchemas);
+
+        assertThat(mapping.nodeLabels()).containsExactlyInAnyOrder(NodeLabel.of("A"), NodeLabel.of("B"));
+        assertThat(mapping.propertySchemas(NodeLabel.of("A"), importPropertySchemas)).isEqualTo(Map.of(
+            "foo", PropertySchema.of("foo", ValueType.LONG, DefaultValue.forLong(), PropertyState.TRANSIENT)
+        ));
+        assertThat(mapping.propertySchemas(NodeLabel.of("B"), importPropertySchemas)).isEqualTo(Map.of(
+            "bar", PropertySchema.of("bar", ValueType.DOUBLE, DefaultValue.forDouble(), PropertyState.PERSISTENT)
+        ));
+    }
 }