Gh-3202 GafferPop filter edges by edge labels

Fix federated test
gchq · May 2, 2024 · b2a029d · b2a029d
1 parent c0b950f
commit b2a029d
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 5 deletions.
diff --git a/library/tinkerpop/src/main/java/uk/gov/gchq/gaffer/tinkerpop/GafferPopGraph.java b/library/tinkerpop/src/main/java/uk/gov/gchq/gaffer/tinkerpop/GafferPopGraph.java
@@ -67,12 +67,15 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
@@ -551,12 +554,35 @@ public Iterator<Edge> edges(final Object... elementIds) {
                 .then(new Limit<>(getAllElementsLimit, true))
                 .build();
         } else {
-            getOperation = new GetElements.Builder()
-                    .input(getElementSeeds(Arrays.asList(elementIds)))
-                    .view(new View.Builder()
+            // Partition elementIds by String[] or Object
+            // The assumption is that the String[] items are edge labels
+            // e.g. elementIds = {"1", ['knows']}
+            Map<Boolean, List<Object>> partitions = Arrays.asList(elementIds).stream()
+                .collect(Collectors.partitioningBy(e -> e instanceof String[]));
+
+            // Flatten arrays into set
+            Set<String> edgeLables = new HashSet<>();
+            partitions.get(true).stream()
+                    .map(e -> Arrays.asList((String[]) e))
+                    .forEach(edgeLables::addAll);
+
+            // Assume everything that isn't a String[] is an element seed
+            List<Object> elements = partitions.get(false);
+
+            if (edgeLables.isEmpty()) {
+                // Get all edges
+                getOperation = new GetElements.Builder()
+                        .input(getElementSeeds(elements))
+                        .view(new View.Builder()
                             .edges(graph.getSchema().getEdgeGroups())
                             .build())
-                    .build();
+                        .build();
+            } else {
+                // filter edges by label using a view
+                return edgesWithView(elements, Direction.BOTH, new View.Builder()
+                    .edges(edgeLables)
+                    .build());
+            }
         }
 
         // Run requested chain on the graph
@@ -901,6 +927,16 @@ private List<ElementSeed> getElementSeeds(final Iterable<Object> ids) {
             // Extract source and destination from ID list
             } else if (id instanceof Iterable) {
                 ((Iterable<?>) id).forEach(edgeIdList::add);
+            // } else if (id instanceof Object[]) {
+            //     // This is probably a label .. 'knows'
+            //     // how do we turn that into an EdgeSeed??
+
+
+            //     Object[] arr = (Object[]) id;
+            //     for(Object o : arr) {
+            //         EdgeSeed seed = new EdgeSeed();
+            //         edgeIdList.add(o);
+            //     }
             // Attempt to extract source and destination IDs from a string form of an array/list
             } else if ((id instanceof String) && (((String) id).matches("^\\[.*,.*\\]$"))) {
                 edgeIdList = Arrays.asList(((String) id)

diff --git a/library/tinkerpop/src/test/java/uk/gov/gchq/gaffer/tinkerpop/GafferPopFederatedIT.java b/library/tinkerpop/src/test/java/uk/gov/gchq/gaffer/tinkerpop/GafferPopFederatedIT.java
@@ -141,7 +141,7 @@ void shouldFilterVertexesByPropertyValue() {
 
         // Then
         assertThat(result)
-                .containsExactly(1.0, 0.8);
+                .containsExactly(0.8);
     }
 
     @Test