Improve tests and refactor

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

Author: s1ck

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

@@ -70,23 +70,27 @@ abstract Map<String, PropertySchema> propertySchemas(
         Map<String, PropertySchema> importPropertySchemas
     );
 
-    static NodeLabelTokenToPropertyKeys merge(
+    /**
+     * Computes the union of the two given mappings without
+     * changing the contents of the mappings themselves.
+     */
+    static NodeLabelTokenToPropertyKeys union(
         NodeLabelTokenToPropertyKeys left,
         NodeLabelTokenToPropertyKeys right,
         Map<String, PropertySchema> importPropertySchemas
     ) {
-        var merge = NodeLabelTokenToPropertyKeys.lazy();
+        var union = NodeLabelTokenToPropertyKeys.lazy();
 
         left.nodeLabels().forEach(nodeLabel -> {
             var propertyKeys = left.propertySchemas(nodeLabel, importPropertySchemas).keySet();
-            merge.add(NodeLabelTokens.ofNodeLabel(nodeLabel), propertyKeys);
+            union.add(NodeLabelTokens.ofNodeLabel(nodeLabel), propertyKeys);
         });
         right.nodeLabels().forEach(nodeLabel -> {
             var propertyKeys = right.propertySchemas(nodeLabel, importPropertySchemas).keySet();
-            merge.add(NodeLabelTokens.ofNodeLabel(nodeLabel), propertyKeys);
+            union.add(NodeLabelTokens.ofNodeLabel(nodeLabel), propertyKeys);
         });
 
-        return merge;
+        return union;
     }
 
     private static class Fixed extends NodeLabelTokenToPropertyKeys {

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

@@ -199,17 +199,11 @@ 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 localLabelTokenToPropertyKeys = closeThreadLocalBuilders();
 
         var idMap = this.idMapBuilder.build(labelInformationBuilder, highestNeoId, concurrency);
         var nodeProperties = buildProperties(idMap);
-        var nodeSchema = buildNodeSchema(idMap, nodeLabelTokenToPropertyKeysList, nodeProperties);
+        var nodeSchema = buildNodeSchema(idMap, localLabelTokenToPropertyKeys, nodeProperties);
         var nodePropertyStore = NodePropertyStore.builder().properties(nodeProperties).build();
 
         return ImmutableNodes.builder()
@@ -219,37 +213,53 @@ public Nodes build(long highestNeoId) {
             .build();
     }
 
+    private List<NodeLabelTokenToPropertyKeys> closeThreadLocalBuilders() {
+        // Flush remaining buffer contents
+        this.threadLocalBuilder.forEach(ThreadLocalBuilder::flush);
+        // Collect token to property keys for final union
+        var labelTokenToPropertyKeys = new ArrayList<NodeLabelTokenToPropertyKeys>();
+        this.threadLocalBuilder.forEach(tlb -> labelTokenToPropertyKeys.add(tlb.nodeLabelTokenToPropertyKeys));
+        // Clean up resources held by local builders
+        this.threadLocalBuilder.close();
+
+        return labelTokenToPropertyKeys;
+    }
+
     private NodeSchema buildNodeSchema(
         IdMap idMap,
-        Collection<NodeLabelTokenToPropertyKeys> localNodeLabelTokenToPropertyKeys,
+        Collection<NodeLabelTokenToPropertyKeys> localLabelTokenToPropertyKeys,
         Map<String, NodeProperty> nodeProperties
     ) {
-        var nodeSchema = NodeSchema.empty();
-        var importPropertySchemas = nodeProperties
+
+        // Collect the property schemas from the imported property values.
+        var propertyKeysToSchema = nodeProperties
             .entrySet()
             .stream()
             .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().propertySchema()));
-
-        // consider node labels without properties
-        var nodeLabels = new HashSet<>(idMap.availableNodeLabels());
-        // and also node labels with associated properties
-        localNodeLabelTokenToPropertyKeys.forEach(mapping -> nodeLabels.addAll(mapping.nodeLabels()));
-        // merge into a global mapping
-        var globalNodeLabelTokenToPropertyKeys = localNodeLabelTokenToPropertyKeys
+        // Union the label to property key mappings from each import thread.
+        var globalLabelTokenToPropertyKeys = localLabelTokenToPropertyKeys
             .stream()
             .reduce(
                 NodeLabelTokenToPropertyKeys.lazy(),
-                (left, right) -> NodeLabelTokenToPropertyKeys.merge(left, right, importPropertySchemas)
-            );
-        // and construct final node property schema
-        nodeLabels.forEach(nodeLabel -> {
-            nodeSchema.addLabel(
-                nodeLabel,
-                globalNodeLabelTokenToPropertyKeys.propertySchemas(nodeLabel, importPropertySchemas)
+                (left, right) -> NodeLabelTokenToPropertyKeys.union(left, right, propertyKeysToSchema)
             );
-        });
+        // Collect node labels without properties from the id map
+        // as they are not stored in the above union mapping.
+        var nodeLabels = new HashSet<>(idMap.availableNodeLabels());
+        // Add labels that actually have node properties attached.
+        localLabelTokenToPropertyKeys.forEach(localMapping -> nodeLabels.addAll(localMapping.nodeLabels()));
 
-        return nodeSchema;
+        // Use all labels and the global label to property
+        // key mapping to construct the final node schema.
+        return nodeLabels.stream()
+            .reduce(
+                NodeSchema.empty(),
+                (unionSchema, nodeLabel) -> unionSchema.addLabel(
+                    nodeLabel,
+                    globalLabelTokenToPropertyKeys.propertySchemas(nodeLabel, propertyKeysToSchema)
+                ),
+                (lhs, rhs) -> lhs
+            );
     }
 
     private Map<String, NodeProperty> buildProperties(IdMap idMap) {
@@ -344,10 +354,6 @@ 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

@@ -204,24 +204,40 @@ void testNodeLabelsFixed() {
     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)
+            "bar", PropertySchema.of("bar", ValueType.DOUBLE, DefaultValue.forDouble(), PropertyState.PERSISTENT),
+            "baz", PropertySchema.of("baz", ValueType.LONG_ARRAY, DefaultValue.forLongArray(), PropertyState.PERSISTENT)
         );
 
-        var left = NodeLabelTokenToPropertyKeys.lazy();
-        left.add(NodeLabelTokens.ofStrings("A"), List.of("foo"));
+        var aLabel = NodeLabel.of("A");
+        var bLabel = NodeLabel.of("B");
+        var cLabel = NodeLabel.of("C");
 
-        var right = NodeLabelTokenToPropertyKeys.lazy();
-        right.add(NodeLabelTokens.ofStrings("B"), List.of("bar"));
+        var mapping0 = NodeLabelTokenToPropertyKeys.lazy();
+        mapping0.add(NodeLabelTokens.ofNodeLabels(aLabel), List.of("foo"));
+        mapping0.add(NodeLabelTokens.ofNodeLabels(cLabel), List.of("bar"));
 
-        var mapping = NodeLabelTokenToPropertyKeys.merge(left, right, importPropertySchemas);
+        var mapping1 = NodeLabelTokenToPropertyKeys.lazy();
+        mapping1.add(NodeLabelTokens.ofNodeLabels(bLabel), List.of("bar"));
 
-        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)
+        var mapping2 = NodeLabelTokenToPropertyKeys.lazy();
+        mapping2.add(NodeLabelTokens.ofNodeLabels(aLabel, cLabel), List.of("baz"));
+
+        var union = NodeLabelTokenToPropertyKeys.union(mapping0, mapping1, importPropertySchemas);
+        union = NodeLabelTokenToPropertyKeys.union(union, mapping2, importPropertySchemas);
+
+        assertThat(union.nodeLabels()).containsExactlyInAnyOrder(aLabel, bLabel, cLabel);
+
+        assertThat(union.propertySchemas(aLabel, importPropertySchemas)).isEqualTo(Map.of(
+            "foo", PropertySchema.of("foo", ValueType.LONG, DefaultValue.forLong(), PropertyState.TRANSIENT),
+            "baz", PropertySchema.of("baz", ValueType.LONG_ARRAY, DefaultValue.forLongArray(), PropertyState.PERSISTENT)
         ));
-        assertThat(mapping.propertySchemas(NodeLabel.of("B"), importPropertySchemas)).isEqualTo(Map.of(
+        assertThat(union.propertySchemas(bLabel, importPropertySchemas)).isEqualTo(Map.of(
             "bar", PropertySchema.of("bar", ValueType.DOUBLE, DefaultValue.forDouble(), PropertyState.PERSISTENT)
         ));
+        assertThat(union.propertySchemas(cLabel, importPropertySchemas)).isEqualTo(Map.of(
+            "bar", PropertySchema.of("bar", ValueType.DOUBLE, DefaultValue.forDouble(), PropertyState.PERSISTENT),
+            "baz", PropertySchema.of("baz", ValueType.LONG_ARRAY, DefaultValue.forLongArray(), PropertyState.PERSISTENT)
+        ));
     }
 }