flatmap query utility

onthegomap · May 24, 2024 · de0418b · de0418b
1 parent 5b4e338
commit de0418b
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 4 deletions.
diff --git a/planetiler-core/src/main/java/com/onthegomap/planetiler/reader/Struct.java b/planetiler-core/src/main/java/com/onthegomap/planetiler/reader/Struct.java
@@ -14,6 +14,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
 
 /**
@@ -240,6 +241,20 @@ default String asJson() {
     return JsonConversion.writeValueAsString(rawValue());
   }
 
+  /**
+   * Returns a new list where each element of this list has been expanded to the list of elements returned by
+   * {@code mapper}.
+   * <p>
+   * Individual items are treated as a list containing just that item.
+   */
+  default Struct flatMap(UnaryOperator<Struct> mapper) {
+    var list = asList().stream()
+      .flatMap(item -> mapper.apply(item).asList().stream())
+      .map(Struct::of)
+      .toList();
+    return list.isEmpty() ? NULL : new ListStruct(list);
+  }
+
   class PrimitiveStruct<T> implements Struct {
 
     final T value;
@@ -513,10 +528,15 @@ public Struct get(Object key) {
       var result = value.get(key);
       if (result != null) {
         return result;
-      } else if (key instanceof String s && s.contains(".")) {
-        String[] parts = s.split("\\.", 2);
-        if (parts.length == 2) {
-          return get(parts[0], parts[1]);
+      } else if (key instanceof String s) {
+        if (s.contains(".")) {
+          String[] parts = s.split("\\.", 2);
+          if (parts.length == 2) {
+            String firstPart = parts[0];
+            return firstPart.endsWith("[]") ?
+              get(firstPart.substring(0, firstPart.length() - 2)).flatMap(child -> child.get(parts[1])) :
+              get(firstPart, parts[1]);
+          }
         }
       }
       return NULL;

diff --git a/planetiler-core/src/test/java/com/onthegomap/planetiler/reader/StructTest.java b/planetiler-core/src/test/java/com/onthegomap/planetiler/reader/StructTest.java
@@ -69,6 +69,16 @@ void testGetDottedFromWithTags() {
     )).hasTag("a.b"));
   }
 
+
+  @Test
+  void testListQuery() {
+    var struct = Struct.of(Map.of(
+      "a", List.of(Map.of("b", "c"), Map.of("b", "d"))
+    ));
+    assertEquals("d", struct.get("a").flatMap(elem -> elem.get("b")).get(1).asString());
+    assertEquals(Struct.of(List.of("c", "d")), struct.get("a[].b"));
+  }
+
   @Test
   void testListGet() {
     var struct = Struct.of(List.of(1, 2, 3));