Working implementation of reverseResolve (single per resource).

aehrc · Oct 19, 2023 · 982cd7f · 982cd7f
1 parent f879e9e
commit 982cd7f
Show file tree

Hide file tree

Showing 33 changed files with 346 additions and 190 deletions.
diff --git a/fhir-server/src/main/java/au/csiro/pathling/security/ga4gh/PassportScopeEnforcer.java b/fhir-server/src/main/java/au/csiro/pathling/security/ga4gh/PassportScopeEnforcer.java
@@ -84,7 +84,7 @@ public Dataset<Row> enforce(@Nonnull final ResourceType subjectResource,
 
       // Build a new expression parser, and parse all the column expressions within the query.
       final ResourceCollection inputContext = ResourceCollection
-          .build(getFhirContext(), getDataSource().read(subjectResource), subjectResource
+          .build(getFhirContext(), subjectResource
           );
       return filterDataset(inputContext, filters, dataset, dataset.col("id"), Column::or);
     }

diff --git a/fhir-server/src/test/java/au/csiro/pathling/extract/ExtractAggregatePOCViewTest.java b/fhir-server/src/test/java/au/csiro/pathling/extract/ExtractAggregatePOCViewTest.java
@@ -45,7 +45,7 @@
 import au.csiro.pathling.view.FromSelection;
 import au.csiro.pathling.view.PrimitiveSelection;
 import au.csiro.pathling.view.Selection;
-import au.csiro.pathling.view.ViewContext;
+import au.csiro.pathling.view.ExecutionContext;
 import ca.uhn.fhir.context.FhirContext;
 import ca.uhn.fhir.parser.IParser;
 import com.google.common.collect.Streams;
@@ -176,8 +176,8 @@ void setUp() {
   }
 
 
-  ViewContext newContext() {
-    return new ViewContext(spark, fhirContext, dataSource);
+  ExecutionContext newContext() {
+    return new ExecutionContext(spark, fhirContext, dataSource);
   }
 
   @Test
@@ -222,36 +222,7 @@ void testExtractView() {
     System.out.println(resultDataset.queryExecution().optimizedPlan());
 
   }
-
-
-  @Test
-  void testExtractWithResoutions() {
-
-    final List<String> expressions = List.of(
-        "id",
-        "gender",
-        "reverseResolve(Condition.subject).clinicalStatus"
-    );
-
-    System.out.println("### Expressions: ###");
-    expressions.forEach(System.out::println);
-    final ExtractRequest extractRequest = ExtractRequest.fromUserInput(
-        ResourceType.PATIENT,
-        Optional.of(expressions),
-        Optional.empty(),
-        Optional.empty()
-    );
-
-    final QueryParser queryParser = new QueryParser(new Parser());
-    final ExtractView extractView = queryParser.toView(extractRequest);
-    System.out.println("## Extract view ##");
-    extractView.printTree();
-    final Dataset<Row> resultDataset = extractView.evaluate(newContext());
-    resultDataset.show(false);
-    System.out.println(resultDataset.logicalPlan());
-    System.out.println(resultDataset.queryExecution().executedPlan());
-  }
-
+
   @Test
   void testAggregation() {
 
@@ -417,7 +388,7 @@ void testAggregationFull() {
     groupByCollections.forEach(c -> System.out.println(c.getFhirType()));
 
     final DatasetResult<Collection> datasetCollectionResult = groupByCollections.stream()
-        .map(k -> (DatasetResult<Collection>) DatasetResult.of(k))
+        .map(k -> (DatasetResult<Collection>) DatasetResult.pureOne(k))
         .reduce(DatasetResult.empty(), DatasetResult::andThen);
 
     System.out.println("### Grouping Collection Result: ###");
@@ -467,4 +438,34 @@ void mockResource(final ResourceType... resourceTypes) {
     TestHelpers.mockResource(dataSource, spark, resourceTypes);
   }
 
+
+
+  @Test
+  void testExtractViewWithReverseResolve() {
+
+    final List<String> expressions = List.of(
+        "id",
+        "gender",
+        "reverseResolve(Condition.subject).clinicalStatus",
+        "reverseResolve(Observation.subject).id.count()"
+    );
+
+    System.out.println("### Expressions: ###");
+    expressions.forEach(System.out::println);
+    final ExtractRequest extractRequest = ExtractRequest.fromUserInput(
+        ResourceType.PATIENT,
+        Optional.of(expressions),
+        Optional.empty(),
+        Optional.empty()
+    );
+
+    final QueryParser queryParser = new QueryParser(new Parser());
+    final ExtractView extractView = queryParser.toView(extractRequest);
+    System.out.println("## Extract view ##");
+    extractView.printTree();
+    final Dataset<Row> resultDataset = extractView.evaluate(newContext());
+    resultDataset.show(false);
+    System.out.println(resultDataset.logicalPlan());
+    System.out.println(resultDataset.queryExecution().executedPlan());
+  }
 }
diff --git a/fhir-server/src/test/java/au/csiro/pathling/fhirpath/parser/AbstractParserTest.java b/fhir-server/src/test/java/au/csiro/pathling/fhirpath/parser/AbstractParserTest.java
@@ -85,12 +85,13 @@ void setUp() {
   @SuppressWarnings("SameParameterValue")
   void setSubjectResource(@Nonnull final ResourceType resourceType) {
     final ResourceCollection subjectResource = ResourceCollection
-        .build(fhirContext, dataSource.read(resourceType), resourceType);
+        .build(fhirContext, resourceType);
 
     evaluationContext = new EvaluationContextBuilder(spark, fhirContext)
         .dataset(dataSource.read(resourceType))
         .resource(subjectResource)
         .inputContext(subjectResource)
+        .dataSource(dataSource)
         .build();
   }
 
@@ -106,7 +107,7 @@ void mockResource(final ResourceType... resourceTypes) {
   protected FhirPathAssertion assertThatResultOf(@Nonnull final ResourceType resourceType,
       @Nonnull final String expression) {
     final ResourceCollection subjectResource = ResourceCollection
-        .build(fhirContext, dataSource.read(resourceType), resourceType);
+        .build(fhirContext, resourceType);
 
     final EvaluationContext evaluationContext = new EvaluationContextBuilder(spark, fhirContext)
         .dataset(dataSource.read(resourceType))

diff --git a/fhir-server/src/test/java/au/csiro/pathling/fhirpath/parser/ParserTest.java b/fhir-server/src/test/java/au/csiro/pathling/fhirpath/parser/ParserTest.java
@@ -45,6 +45,7 @@
 import au.csiro.pathling.test.builders.DatasetBuilder;
 import au.csiro.pathling.test.helpers.TerminologyServiceHelpers;
 import au.csiro.pathling.test.helpers.TerminologyServiceHelpers.TranslateExpectations;
+import org.apache.spark.sql.execution.SparkPlan;
 import org.apache.spark.sql.types.DataTypes;
 import org.hl7.fhir.r4.model.Enumerations.ResourceType;
 import org.junit.jupiter.api.Assertions;
@@ -57,7 +58,7 @@ public class ParserTest extends AbstractParserTest {
 
   @SuppressWarnings("SameParameterValue")
   private <T extends Throwable> T assertThrows(final Class<T> errorType, final String expression) {
-    return Assertions.assertThrows(errorType, () -> parser.evaluate(expression, evaluationContext));
+    return Assertions.assertThrows(errorType, () -> parser.evaluatePure(expression, evaluationContext));
   }
 
   private TranslateExpectations setupMockTranslationFor_195662009_444814009(
@@ -245,6 +246,23 @@ void testCodingOperations() {
   //       .hasCodingValue(expectedCodingWithVersion);
   // }
 
+
+  @Test
+  void testBasicReverseResolve() {
+    mockResource(ResourceType.PATIENT, ResourceType.CONDITION);
+
+    final SparkPlan p = assertThatResultOf(
+        "reverseResolve(Condition.subject).code.coding")
+        //.isElementPath(IntegerCollection.class)
+        //.isSingular()
+        .selectOrderedResult()
+        .debugAllRows()
+        .getDataset().queryExecution().executedPlan();
+
+    System.out.println(p);
+  }
+
+
   @Test
   void testCountWithReverseResolve() {
     assertThatResultOf("reverseResolve(Condition.subject).code.coding.count()")

diff --git a/fhirpath/src/main/java/au/csiro/pathling/aggregate/AggregateQueryExecutor.java b/fhirpath/src/main/java/au/csiro/pathling/aggregate/AggregateQueryExecutor.java
@@ -27,7 +27,7 @@
 import au.csiro.pathling.query.QueryParser;
 import au.csiro.pathling.terminology.TerminologyServiceFactory;
 import au.csiro.pathling.view.AggregationView;
-import au.csiro.pathling.view.ViewContext;
+import au.csiro.pathling.view.ExecutionContext;
 import ca.uhn.fhir.context.FhirContext;
 import java.util.Collections;
 import java.util.List;
@@ -81,8 +81,8 @@ public ResultWithExpressions buildQuery(@Nonnull final AggregateRequest query) {
         Collections.emptyList(), Collections.emptyList());
   }
 
-  protected ViewContext newContext() {
-    return new ViewContext(sparkSession, fhirContext, dataSource);
+  protected ExecutionContext newContext() {
+    return new ExecutionContext(sparkSession, fhirContext, dataSource);
   }
 
   //

diff --git a/fhirpath/src/main/java/au/csiro/pathling/extract/ExtractQueryExecutor.java b/fhirpath/src/main/java/au/csiro/pathling/extract/ExtractQueryExecutor.java
@@ -8,7 +8,7 @@
 import au.csiro.pathling.query.QueryParser;
 import au.csiro.pathling.terminology.TerminologyServiceFactory;
 import au.csiro.pathling.view.ExtractView;
-import au.csiro.pathling.view.ViewContext;
+import au.csiro.pathling.view.ExecutionContext;
 import ca.uhn.fhir.context.FhirContext;
 import java.util.Optional;
 import javax.annotation.Nonnull;
@@ -64,8 +64,8 @@ public Dataset<Row> buildQuery(@Nonnull final ExtractRequest query,
     return query.getLimit().map(resultDataset::limit).orElse(resultDataset);
   }
 
-  protected ViewContext newContext() {
-    return new ViewContext(sparkSession, fhirContext, dataSource);
+  protected ExecutionContext newContext() {
+    return new ExecutionContext(sparkSession, fhirContext, dataSource);
   }
 
 }
diff --git a/fhirpath/src/main/java/au/csiro/pathling/fhirpath/FhirPath.java b/fhirpath/src/main/java/au/csiro/pathling/fhirpath/FhirPath.java
@@ -72,7 +72,7 @@ class This<I extends Collection> implements FhirPath<I> {
     @Override
     public DatasetResult.One<I> eval(@Nonnull final I input,
         @Nonnull final EvaluationContext context) {
-      return DatasetResult.of(input);
+      return DatasetResult.pureOne(input);
     }
 
     @Override
@@ -105,7 +105,7 @@ class Composite<I extends Collection> implements FhirPath<I> {
     public DatasetResult.One<I> eval(@Nonnull final I input,
         @Nonnull final EvaluationContext context) {
       return elements.stream()
-          .reduce(DatasetResult.of(input),
+          .reduce(DatasetResult.pureOne(input),
               (acc, element) -> element.eval(acc.getValue(), context).withTransformOf(acc),
               (a, b) -> b);
     }

diff --git a/fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/ResourceCollection.java b/fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/ResourceCollection.java
@@ -30,17 +30,12 @@
 import ca.uhn.fhir.context.FhirContext;
 import ca.uhn.fhir.context.RuntimeResourceDefinition;
 import java.util.EnumSet;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.function.Function;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import lombok.Getter;
 import org.apache.spark.sql.Column;
-import org.apache.spark.sql.Dataset;
-import org.apache.spark.sql.Row;
 import org.apache.spark.sql.functions;
 import org.hl7.fhir.exceptions.FHIRException;
 import org.hl7.fhir.r4.model.Enumerations.FHIRDefinedType;
@@ -55,35 +50,24 @@
 @Getter
 public class ResourceCollection extends Collection {
 
-  /**
-   * A mapping between the names of elements in the resource and the corresponding {@link Column}.
-   */
-  @Nonnull
-  private final Map<String, Column> elementsToColumns;
-
   /**
    * The {@link ResourceDefinition} for this resource type.
    */
   @Nonnull
   private final ResourceDefinition resourceDefinition;
 
-  /**
-   * The {@link Dataset} containing the resource data.
-   */
   @Nonnull
-  private final Dataset<Row> dataset;
+  final Optional<String> prefix;
 
   protected ResourceCollection(@Nonnull final Column column,
       @Nonnull final Optional<FhirPathType> type,
       @Nonnull final Optional<FHIRDefinedType> fhirType,
       @Nonnull final Optional<? extends NodeDefinition> definition,
-      @Nonnull final Map<String, Column> elementsToColumns,
       @Nonnull final ResourceDefinition resourceDefinition,
-      @Nonnull final Dataset<Row> dataset) {
+      @Nonnull final Optional<String> prefix) {
     super(column, type, fhirType, definition);
-    this.elementsToColumns = elementsToColumns;
     this.resourceDefinition = resourceDefinition;
-    this.dataset = dataset;
+    this.prefix = prefix;
   }
 
   @Nonnull
@@ -99,29 +83,30 @@ private static Optional<FHIRDefinedType> getFhirType(@Nonnull final ResourceType
    * Build a new ResourcePath using the supplied {@link FhirContext} and {@link DataSource}.
    *
    * @param fhirContext the {@link FhirContext} to use for sourcing the resource definition
-   * @param dataset the {@link Dataset} that contains the resource data
    * @param resourceType the type of the resource
    * @return A shiny new ResourcePath
    */
   @Nonnull
   public static ResourceCollection build(@Nonnull final FhirContext fhirContext,
-      @Nonnull final Dataset<Row> dataset, @Nonnull final ResourceType resourceType) {
+      @Nonnull final ResourceType resourceType, @Nonnull final Optional<String> prefix) {
 
     // Get the resource definition from HAPI.
     final String resourceCode = resourceType.toCode();
     final RuntimeResourceDefinition hapiDefinition = fhirContext.getResourceDefinition(
         resourceCode);
     final ResourceDefinition definition = new ResourceDefinition(resourceType, hapiDefinition);
 
-    //noinspection ReturnOfNull
-    final Map<String, Column> elementsToColumns = Stream.of(dataset.columns())
-        .collect(Collectors.toUnmodifiableMap(Function.identity(), dataset::col));
-
     // We use a literal column as the resource value - the actual value is not important.
     // But the non-null value indicates that the resource should be included in any result.
     return new ResourceCollection(functions.lit(true), Optional.empty(),
-        getFhirType(resourceType), Optional.of(definition), elementsToColumns, definition,
-        dataset);
+        getFhirType(resourceType), Optional.of(definition), definition, prefix);
+  }
+
+
+  @Nonnull
+  public static ResourceCollection build(@Nonnull final FhirContext fhirContext,
+      @Nonnull final ResourceType resourceType) {
+    return build(fhirContext, resourceType, Optional.empty());
   }
 
   /**
@@ -149,7 +134,14 @@ public static Set<ResourceType> supportedResourceTypes() {
    */
   @Nonnull
   public Optional<Column> getElementColumn(@Nonnull final String elementName) {
-    return Optional.ofNullable(elementsToColumns.get(elementName));
+    // TODO: do more checking or not (either based on schema or definition)
+    // return resourceDefinition
+    //     .getChildElement(elementName)
+    //     .map(ChildDefinition::getName)
+    //     .map(functions::col);
+    return prefix
+        .map(p -> p + "_" + elementName).or(() -> Optional.of(elementName))
+        .map(functions::col);
   }
 
   @Nonnull
@@ -178,13 +170,12 @@ protected Collection traverseElement(@Nonnull final ElementDefinition childDef)
             functions.when(getCtx().getValue().isNotNull(), value),
             childDef)).get();
   }
-
-
+
   @Nonnull
   @Override
   public Collection copyWith(@Nonnull final Column newValue) {
     return new ResourceCollection(newValue, getType(), getFhirType(), getDefinition(),
-        elementsToColumns, resourceDefinition, dataset);
+        resourceDefinition, prefix);
   }
 
   @Nonnull

diff --git a/fhirpath/src/main/java/au/csiro/pathling/fhirpath/function/StandardFunctions.java b/fhirpath/src/main/java/au/csiro/pathling/fhirpath/function/StandardFunctions.java
@@ -252,7 +252,6 @@ public static ResourceCollection reverseResolve(@Nonnull final Collection input,
 
     final ResourceCollection foreignResource = ResourceCollection.build(
         evaluationContext.getFhirContext(),
-        dataset,
         foreignResourceType);
 
     final Column referenceColumn = foreignResource.traverse(
@@ -271,7 +270,7 @@ public static ResourceCollection reverseResolve(@Nonnull final Collection input,
     final Dataset<Row> joinedDataset = dataset.join(foreignDataset,
         dataset.col("id_versioned").equalTo(foreignDataset.col("ref_xxx")), "left_outer");
 
-    return ResourceCollection.build(evaluationContext.getFhirContext(), joinedDataset,
+    return ResourceCollection.build(evaluationContext.getFhirContext(),
         foreignResourceType);
 
   }

diff --git a/fhirpath/src/main/java/au/csiro/pathling/fhirpath/parser/InvocationVisitor.java b/fhirpath/src/main/java/au/csiro/pathling/fhirpath/parser/InvocationVisitor.java
@@ -40,6 +40,7 @@
 import au.csiro.pathling.fhirpath.path.Paths.EvalFunction;
 import au.csiro.pathling.fhirpath.path.Paths.This;
 import au.csiro.pathling.fhirpath.path.Paths.Traversal;
+import au.csiro.pathling.fhirpath.path.ReverseResolvePath;
 
 /**
  * This class is invoked on the right-hand side of the invocation expression, and can optionally be
@@ -94,7 +95,11 @@ public FhirPath<Collection> visitFunctionInvocation(
             .collect(toList())
         ).orElse(Collections.emptyList());
 
-    return new EvalFunction(functionIdentifier, arguments);
+    if ("reverseResolve".equals(functionIdentifier)) {
+      return ReverseResolvePath.from(arguments);
+    } else {
+      return new EvalFunction(functionIdentifier, arguments);
+    }
   }
 
   @Override