Add NodeLabelToken#nodeLabels convenience method

Co-authored-by: Paul Horn <paul.horn@neotechnology.com>

Author: s1ck

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

@@ -22,6 +22,9 @@
 import org.jetbrains.annotations.NotNull;
 import org.neo4j.gds.NodeLabel;
 
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
 public interface NodeLabelToken {
 
     /**
@@ -46,6 +49,13 @@ public interface NodeLabelToken {
     @NotNull NodeLabel get(int index);
 
     String[] getStrings();
+
+    /**
+     * @return a stream of {@link org.neo4j.gds.NodeLabel}s represented by this token.
+     */
+    default Stream<NodeLabel> nodeLabels() {
+        return IntStream.range(0, this.size()).mapToObj(this::get);
+    }
 }
 
 

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

@@ -29,7 +29,6 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 abstract class NodeLabelTokenToPropertyKeys {
 
@@ -158,7 +157,7 @@ Set<NodeLabel> nodeLabels() {
                 .keySet()
                 .stream()
                 .map(nodeLabelToken -> nodeLabelToken.isEmpty() ? NodeLabelTokens.ofNodeLabel(NodeLabel.ALL_NODES) : nodeLabelToken)
-                .flatMap(nodeLabelToken -> IntStream.range(0, nodeLabelToken.size()).mapToObj(nodeLabelToken::get))
+                .flatMap(NodeLabelToken::nodeLabels)
                 .collect(Collectors.toSet());
         }
 

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

@@ -0,0 +1,45 @@
+/*
+ * 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.core.loading.construction;
+
+import org.junit.jupiter.api.Test;
+import org.neo4j.gds.NodeLabel;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class NodeLabelTokensTest {
+
+    @Test
+    void testStream() {
+        var aLabel = NodeLabel.of("A");
+        var bLabel = NodeLabel.of("B");
+        var cLabel = NodeLabel.of("C");
+
+        assertThat(NodeLabelTokens.empty().nodeLabels()).isEmpty();
+        assertThat(NodeLabelTokens.ofNodeLabel(aLabel).nodeLabels()).contains(aLabel);
+        assertThat(
+            NodeLabelTokens.ofNodeLabels(aLabel, bLabel, cLabel).nodeLabels()
+        ).contains(aLabel, bLabel, cLabel);
+        assertThat(
+            NodeLabelTokens.ofStrings(aLabel.name(), bLabel.name(), cLabel.name()).nodeLabels()
+        ).contains(aLabel, bLabel, cLabel);
+    }
+
+}