Extracting implicit unnesting functionality for 'extract()' to Implic…

…itUnnester and fixing issues with unnesing of aggregate functions.

fhirpath/src/main/java/au/csiro/pathling/extract/ExtractQueryExecutor.java

@@ -1,15 +1,14 @@
 package au.csiro.pathling.extract;
 
 import static au.csiro.pathling.utilities.Strings.randomAlias;
-import static java.util.stream.Collectors.groupingBy;
-import static java.util.stream.Collectors.mapping;
 import static java.util.stream.Collectors.toList;
 
 import au.csiro.pathling.QueryExecutor;
 import au.csiro.pathling.config.QueryConfiguration;
 import au.csiro.pathling.fhirpath.FhirPath;
 import au.csiro.pathling.fhirpath.execution.MultiFhirpathEvaluator;
 import au.csiro.pathling.fhirpath.parser.Parser;
+import au.csiro.pathling.fhirpath.path.Paths.This;
 import au.csiro.pathling.io.source.DataSource;
 import au.csiro.pathling.query.ExpressionWithLabel;
 import au.csiro.pathling.terminology.TerminologyServiceFactory;
@@ -24,9 +23,7 @@
 import jakarta.annotation.Nonnull;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.LinkedHashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Stream;
 import lombok.extern.slf4j.Slf4j;
@@ -158,33 +155,35 @@ private Projection buildProjection(@Nonnull final ExtractRequest query,
         constraint);
   }
 
+  static UnnestingSelection fromTree(@Nonnull final Tree<FhirPath> tree) {
+    if (tree instanceof Tree.Leaf<FhirPath> leaf) {
+      // for each leaf we create an unnesting selection with $this as the path
+      return new UnnestingSelection(leaf.getValue(), List.of(
+          new ColumnSelection(
+              List.of(
+                  new RequestedColumn(new This(), randomAlias(), false, Optional.empty())
+              ))
+      ), true);
+    } else if (tree instanceof Tree.Node<FhirPath> node) {
+      // each node represents an unnesting selection of its children
+      return new UnnestingSelection(node.getValue(), node.getChildren().stream()
+          .map(ExtractQueryExecutor::fromTree)
+          .collect(toList()), true);
+    } else {
+      throw new IllegalArgumentException("Unknown tree type: " + tree.getClass());
+    }
+  }
+
   @Nonnull
   static ProjectionClause buildSelectClause(@Nonnull final List<FhirPath> paths) {
     if (paths.isEmpty()) {
       throw new IllegalArgumentException("Empty column list");
     }
-
-    // If there is only one path, and it is a reference to $this, return a simple column selection.
-    // This is the terminal state of the recursion.
-    if (paths.size() == 1 && paths.get(0).isNull()) {
-      final RequestedColumn requestedColumn = new RequestedColumn(paths.get(0),
-          randomAlias(), false, Optional.empty());
-      return new ColumnSelection(List.of(requestedColumn));
-    }
-
-    // Group the paths by their first element. We use a LinkedHashMap to preserve the order.
-    final Map<FhirPath, List<FhirPath>> groupedPaths = paths.stream()
-        .collect(
-            groupingBy(FhirPath::first, LinkedHashMap::new, mapping(FhirPath::suffix, toList())));
-
-    final List<ProjectionClause> selects = groupedPaths.entrySet().stream()
-        .map(entry -> {
-          // we need to split the suffixed by aggregated and non aggregated
-          ProjectionClause tail = buildSelectClause(entry.getValue());
-          return new UnnestingSelection(entry.getKey(), List.of(tail), true);
-        })
-        .collect(toList());
-
+    final Tree<FhirPath> unnestingTree = new ImplicitUnnester().unnestPaths(paths);
+    log.debug("Unnested tree:\n{}", unnestingTree.map(FhirPath::toExpression).toTreeString());
+    // this should be a selection for the resource with the unnesting tree
+    final UnnestingSelection resourceSelection = fromTree(unnestingTree);
+    final List<ProjectionClause> selects = resourceSelection.getComponents();
     // If there is more than one select, return a GroupingSelection.
     // Otherwise, return the single select by itself.
     return selects.size() > 1

fhirpath/src/main/java/au/csiro/pathling/extract/ImplicitUnnester.java

@@ -41,6 +41,19 @@ static List<String> asExpressions(@Nonnull final List<FhirPath> paths) {
     return paths.stream().map(FhirPath::toExpression).toList();
   }
 
+  @Nonnull
+  private static Tree<FhirPath> maybeUnnestingNode(
+      @Nonnull final FhirPath prefix,
+      @Nonnull final List<Tree<FhirPath>> children) {
+    if (children.isEmpty()) {
+      return Tree.Leaf.of(prefix);
+    } else if (children.size() == 1) {
+      return children.get(0).mapValue(prefix::andThen);
+    } else {
+      return Tree.node(prefix, children);
+    }
+  }
+
   @Nonnull
   List<Tree<FhirPath>> unnestPathsInternal(@Nonnull final List<FhirPath> paths) {
     log.trace("Unnesting paths: {}", asExpressions(paths));
@@ -72,7 +85,8 @@ List<Tree<FhirPath>> unnestPathsInternal(@Nonnull final List<FhirPath> paths) {
                 final Stream<Tree<FhirPath>> unnestedNodesStream =
                     suffixesToUnnest.isEmpty()
                     ? Stream.empty()
-                    : Stream.of(Tree.node(entry.getKey(), unnestPathsInternal(suffixesToUnnest)));
+                    : Stream.of(maybeUnnestingNode(entry.getKey(),
+                        unnestPathsInternal(suffixesToUnnest)));
                 final List<Tree<FhirPath>> unnestNodes = unnestedNodesStream.toList();
                 return Stream.concat(unnestNodes.stream(), aggNodes.stream());
               }

fhirpath/src/test/java/au/csiro/pathling/extract/ImplicitUnnesterTest.java

@@ -9,8 +9,6 @@
 import java.util.List;
 import org.junit.jupiter.api.Test;
 
-// TODO: Maybe we should build a nesting three structure first 
-//  and only then convert it to a projection clause
 public class ImplicitUnnesterTest {
 
 
@@ -53,11 +51,9 @@ void testNestedSimpleColumns() {
                 node("given"),
                 node("family")
             ),
-            node("maritalStatus",
-                node("coding",
-                    node("system"),
-                    node("code")
-                )
+            node("maritalStatus.coding",
+                node("system"),
+                node("code")
             )
         ),
         projection
@@ -95,6 +91,51 @@ void testNestedSimpleAndAggregateColumns() {
     );
   }
 
+
+  @Test
+  void testNonBranchedPathsShouldNotBeUnnested() {
+    final List<String> columns = List.of(
+        "id",
+        "name.given.count()"
+    );
+    final Tree<FhirPath> projection = toProjection(columns);
+    // this should be 
+    // id
+    // name
+    //   given
+    //   family
+    // name.count()
+    // name.exists()
+    assertTreeEquals(
+        node("%resource",
+            node("id"),
+            node("name.given.count()")
+        ),
+        projection
+    );
+  }
+
+  @Test
+  void testMixedNestedAggAndValue() {
+    final List<String> columns = List.of(
+        "id",
+        "name.family",
+        "name.given.count()"
+    );
+    final Tree<FhirPath> projection = toProjection(columns);
+    assertTreeEquals(
+        node("%resource",
+            node("id"),
+            node("name",
+                node("family"),
+                node("given.count()")
+            )
+        ),
+        projection
+    );
+  }
+
+
   @Test
   void testNestedSimpleAndNestedAggregateColumns() {
     final List<String> columns = List.of(