Refactor implementation of getResourceKey and getReferenceKey

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/BooleanCollection.java

@@ -23,7 +23,6 @@
 import au.csiro.pathling.fhirpath.column.ColumnRepresentation;
 import au.csiro.pathling.fhirpath.column.DefaultRepresentation;
 import au.csiro.pathling.fhirpath.definition.NodeDefinition;
-import com.google.common.collect.ImmutableSet;
 import java.util.Optional;
 import javax.annotation.Nonnull;
 import org.apache.spark.sql.Row;
@@ -38,9 +37,6 @@
 public class BooleanCollection extends Collection implements Materializable<BooleanType>,
     Comparable {
 
-  private static final ImmutableSet<Class<? extends Comparable>> COMPARABLE_TYPES = ImmutableSet
-      .of(BooleanCollection.class);
-
   protected BooleanCollection(@Nonnull final ColumnRepresentation columnRepresentation,
       @Nonnull final Optional<FhirPathType> type,
       @Nonnull final Optional<FHIRDefinedType> fhirType,
@@ -121,7 +117,7 @@ public Optional<BooleanType> getFhirValueFromRow(@Nonnull final Row row, final i
 
   @Override
   public boolean isComparableTo(@Nonnull final Collection path) {
-    return COMPARABLE_TYPES.contains(path.getClass());
+    return path instanceof BooleanCollection || super.isComparableTo(path);
   }
 
 }

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/CodingCollection.java

@@ -135,6 +135,11 @@ public Optional<Coding> getFhirValueFromRow(@Nonnull final Row row, final int co
     return Optional.of(coding);
   }
 
+  @Override
+  public boolean isComparableTo(@Nonnull final Collection path) {
+    return path instanceof CodingCollection || super.isComparableTo(path);
+  }
+
   @Override
   @Nonnull
   public Function<Comparable, Column> getComparison(@Nonnull final ComparisonOperation operation) {

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/Collection.java

@@ -28,12 +28,12 @@
 import au.csiro.pathling.fhirpath.collection.mixed.MixedCollection;
 import au.csiro.pathling.fhirpath.column.ColumnRepresentation;
 import au.csiro.pathling.fhirpath.column.DefaultRepresentation;
-import au.csiro.pathling.fhirpath.column.NullRepresentation;
 import au.csiro.pathling.fhirpath.definition.ChildDefinition;
 import au.csiro.pathling.fhirpath.definition.ChoiceChildDefinition;
 import au.csiro.pathling.fhirpath.definition.ElementChildDefinition;
 import au.csiro.pathling.fhirpath.definition.ElementDefinition;
 import au.csiro.pathling.fhirpath.definition.NodeDefinition;
+import au.csiro.pathling.fhirpath.function.ColumnTransform;
 import com.google.common.collect.ImmutableMap;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
@@ -82,6 +82,7 @@ public class Collection implements Comparable, Numeric {
           .put(FHIRDefinedType.CODING, CodingCollection.class)
           .put(FHIRDefinedType.QUANTITY, QuantityCollection.class)
           .put(FHIRDefinedType.SIMPLEQUANTITY, QuantityCollection.class)
+          .put(FHIRDefinedType.REFERENCE, ReferenceCollection.class)
           .build();
 
   /**
@@ -191,8 +192,8 @@ private static Collection getInstance(@Nonnull final ColumnRepresentation column
     final FHIRDefinedType resolvedType = fhirType
         .or(() -> definition.flatMap(ElementDefinition::getFhirType))
         .orElseThrow(() -> new IllegalArgumentException("Must have a fhirType or a definition"));
-    final Class<? extends Collection> elementPathClass = classForType(resolvedType).orElse(
-        Collection.class);
+    final Class<? extends Collection> elementPathClass = classForType(resolvedType)
+        .orElse(Collection.class);
     final Optional<FhirPathType> fhirPathType = FhirPathType.forFhirType(resolvedType);
 
     try {
@@ -307,7 +308,7 @@ public Function<Comparable, Column> getComparison(@Nonnull final ComparisonOpera
 
   @Override
   public boolean isComparableTo(@Nonnull final Collection path) {
-    return path.getClass().equals(this.getClass()) || path.getClass().equals(Collection.class);
+    return path instanceof EmptyCollection;
   }
 
   @Nonnull
@@ -328,19 +329,6 @@ public Optional<Column> getNumericContext() {
     return Optional.empty();
   }
 
-
-  /**
-   * Creates a null {@link Collection}.
-   *
-   * @return the null collection.
-   */
-  @Nonnull
-  public static Collection nullCollection() {
-    return new Collection(NullRepresentation.getInstance(), Optional.empty(),
-        Optional.of(FHIRDefinedType.NULL),
-        Optional.empty());
-  }
-
   /**
    * Returns a new {@link Collection} with the specified {@link ColumnRepresentation}.
    *
@@ -363,7 +351,7 @@ public Collection copyWith(@Nonnull final ColumnRepresentation newValue) {
    */
   @Nonnull
   public Collection filter(
-      @Nonnull final Function<ColumnRepresentation, ColumnRepresentation> lambda) {
+      @Nonnull final ColumnTransform lambda) {
     return map(
         ctx -> ctx.filter(col -> lambda.apply(new DefaultRepresentation(col)).getValue()));
   }
@@ -386,7 +374,7 @@ public Collection asSingular() {
    */
   @Nonnull
   public Collection map(
-      @Nonnull final Function<ColumnRepresentation, ColumnRepresentation> mapper) {
+      @Nonnull final ColumnTransform mapper) {
     return copyWith(mapper.apply(getColumn()));
   }
 
@@ -400,7 +388,7 @@ public Collection map(
    */
   @Nonnull
   public <C extends Collection> C map(
-      @Nonnull final Function<ColumnRepresentation, ColumnRepresentation> mapper,
+      @Nonnull final ColumnTransform mapper,
       @Nonnull final Function<ColumnRepresentation, C> constructor) {
     return constructor.apply(mapper.apply(getColumn()));
   }
@@ -425,7 +413,7 @@ public Column getColumnValue() {
    */
   @Nonnull
   public Collection filterByType(@Nonnull final TypeSpecifier type) {
-    return Collection.nullCollection();
+    return EmptyCollection.getInstance();
   }
 
   // TODO: Remove this after removing usages.

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/DateCollection.java

@@ -117,7 +117,8 @@ public Function<Comparable, Column> getComparison(@Nonnull final ComparisonOpera
 
   @Override
   public boolean isComparableTo(@Nonnull final Collection path) {
-    return DateTimeCollection.getComparableTypes().contains(path.getClass());
+    return DateTimeCollection.getComparableTypes().contains(path.getClass())
+        || super.isComparableTo(path);
   }
 
   @Nonnull

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/DateTimeCollection.java

@@ -159,7 +159,7 @@ public Function<Comparable, Column> getComparison(@Nonnull final ComparisonOpera
 
   @Override
   public boolean isComparableTo(@Nonnull final Collection path) {
-    return COMPARABLE_TYPES.contains(path.getClass());
+    return COMPARABLE_TYPES.contains(path.getClass()) || super.isComparableTo(path);
   }
 
   @Nonnull

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/DecimalCollection.java

@@ -145,7 +145,8 @@ public static org.apache.spark.sql.types.DecimalType getDecimalType() {
 
   @Override
   public boolean isComparableTo(@Nonnull final Collection path) {
-    return IntegerCollection.getComparableTypes().contains(path.getClass());
+    return IntegerCollection.getComparableTypes().contains(path.getClass()) ||
+        super.isComparableTo(path);
   }
 
   @Nonnull

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/EmptyCollection.java

@@ -0,0 +1,37 @@
+package au.csiro.pathling.fhirpath.collection;
+
+import au.csiro.pathling.fhirpath.FhirPathType;
+import au.csiro.pathling.fhirpath.column.ColumnRepresentation;
+import au.csiro.pathling.fhirpath.column.EmptyRepresentation;
+import au.csiro.pathling.fhirpath.definition.NodeDefinition;
+import java.util.Optional;
+import javax.annotation.Nonnull;
+import org.hl7.fhir.r4.model.Enumerations.FHIRDefinedType;
+
+/**
+ * Represents an empty collection.
+ */
+public class EmptyCollection extends Collection {
+
+  private static final EmptyCollection INSTANCE = new EmptyCollection(
+      EmptyRepresentation.getInstance(), Optional.empty(), Optional.empty(), Optional.empty());
+
+  protected EmptyCollection(@Nonnull final ColumnRepresentation column,
+      @Nonnull final Optional<FhirPathType> type, @Nonnull final Optional<FHIRDefinedType> fhirType,
+      @Nonnull final Optional<? extends NodeDefinition> definition) {
+    super(column, type, fhirType, definition);
+  }
+
+  /**
+   * @return A singleton instance of this class
+   */
+  public static EmptyCollection getInstance() {
+    return INSTANCE;
+  }
+
+  @Override
+  public boolean isComparableTo(@Nonnull final Collection path) {
+    return true;
+  }
+
+}

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/IntegerCollection.java

@@ -192,7 +192,7 @@ public static ImmutableSet<Class<? extends Comparable>> getComparableTypes() {
 
   @Override
   public boolean isComparableTo(@Nonnull final Collection path) {
-    return COMPARABLE_TYPES.contains(path.getClass());
+    return COMPARABLE_TYPES.contains(path.getClass()) || super.isComparableTo(path);
   }
 
   @Nonnull

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/QuantityCollection.java

@@ -228,6 +228,11 @@ private static QuantityCollection buildLiteralPath(@Nonnull final BigDecimal dec
         new DefaultRepresentation(QuantityEncoding.encodeLiteral(quantity)));
   }
 
+  @Override
+  public boolean isComparableTo(@Nonnull final Collection path) {
+    return path instanceof QuantityCollection || super.isComparableTo(path);
+  }
+
   @Nonnull
   @Override
   public Function<Comparable, Column> getComparison(@Nonnull final ComparisonOperation operation) {
@@ -268,8 +273,9 @@ public Function<Numeric, Collection> getMathOperation(@Nonnull final MathOperati
       final Column resultStruct = QuantityEncoding.toStruct(
           sourceContext.getField("id"),
           FlexiDecimal.toDecimal(resultColumn),
-          // NOTE: This (setting value_scale to null) works because we never decode this struct to a Quantity.
-          // The only Quantities that are decoded are calendar duration quantities parsed from literals.
+          // NOTE: This (setting value_scale to null) works because we never decode this struct to a 
+          // Quantity. The only Quantities that are decoded are calendar duration quantities parsed 
+          // from literals.
           lit(null),
           sourceContext.getField("comparator"),
           resultCode,

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/ReferenceCollection.java

@@ -0,0 +1,56 @@
+package au.csiro.pathling.fhirpath.collection;
+
+import au.csiro.pathling.fhirpath.FhirPathType;
+import au.csiro.pathling.fhirpath.TypeSpecifier;
+import au.csiro.pathling.fhirpath.column.ColumnRepresentation;
+import au.csiro.pathling.fhirpath.definition.NodeDefinition;
+import au.csiro.pathling.fhirpath.function.ColumnTransform;
+import java.util.Optional;
+import javax.annotation.Nonnull;
+import org.hl7.fhir.r4.model.Enumerations.FHIRDefinedType;
+
+/**
+ * Represents a collection of Reference elements.
+ *
+ * @author John Grimes
+ * @see <a href="https://hl7.org/fhir/R4/references.html#Reference">Resource References</a>
+ */
+public class ReferenceCollection extends Collection {
+
+  private static final String REFERENCE_ELEMENT_NAME = "reference";
+
+  protected ReferenceCollection(@Nonnull final ColumnRepresentation column,
+      @Nonnull final Optional<FhirPathType> type, @Nonnull final Optional<FHIRDefinedType> fhirType,
+      @Nonnull final Optional<? extends NodeDefinition> definition) {
+    super(column, type, fhirType, definition);
+  }
+
+  /**
+   * @param typeSpecifier The type specifier to filter by
+   * @return a {@link Collection} containing the keys of the references in this collection, suitable
+   * for joining with resource keys
+   */
+  @Nonnull
+  public Collection getKeyCollection(@Nonnull final Optional<TypeSpecifier> typeSpecifier) {
+    return typeSpecifier
+        // If a type was specified, create a regular expression that matches references of this type.
+        .map(ts -> ts.toFhirType().toCode() + "/.+")
+        // Get a ColumnTransform that filters the reference column based on the regular expression.
+        .map(this::keyFilter)
+        // Apply the filter to the reference column.
+        .map(this::filter)
+        // Return a StringCollection of the reference elements.s
+        .flatMap(c -> c.traverse(REFERENCE_ELEMENT_NAME))
+        // If no type was specified, return the reference column as is.
+        .or(() -> this.traverse(REFERENCE_ELEMENT_NAME))
+        // If the reference column is not present, return an empty collection.
+        .orElse(EmptyCollection.getInstance());
+  }
+
+  @Nonnull
+  private ColumnTransform keyFilter(@Nonnull final String pattern) {
+    return col -> col.traverse(REFERENCE_ELEMENT_NAME, Optional.of(FHIRDefinedType.STRING))
+        .like(pattern);
+  }
+
+}

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/ResourceCollection.java

@@ -174,8 +174,10 @@ public Collection copyWith(@Nonnull final ColumnRepresentation newValue) {
    * @return A column that can be used as a key for joining to this resource type
    */
   @Nonnull
-  public ColumnRepresentation getKeyColumn() {
-    return getColumn().traverse("id_versioned", Optional.of(FHIRDefinedType.STRING));
+  public Collection getKeyCollection() {
+    return StringCollection.build(
+        getColumn().traverse("id_versioned", Optional.of(FHIRDefinedType.STRING))
+    );
   }
 
 }

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/StringCollection.java

@@ -269,4 +269,10 @@ public Optional<PrimitiveType> getFhirValueFromRow(@Nonnull final Row row,
   public StringCollection asStringPath() {
     return this;
   }
+
+  @Override
+  public boolean isComparableTo(@Nonnull final Collection path) {
+    return path instanceof StringCollection || super.isComparableTo(path);
+  }
+
 }

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/TimeCollection.java

@@ -109,4 +109,9 @@ public StringCollection asStringPath() {
     return map(ColumnRepresentation::asString, StringCollection::build);
   }
 
+  @Override
+  public boolean isComparableTo(@Nonnull final Collection path) {
+    return path instanceof TimeCollection || super.isComparableTo(path);
+  }
+
 }

fhirpath/src/main/java/au/csiro/pathling/fhirpath/collection/mixed/MixedCollection.java

@@ -1,7 +1,7 @@
 package au.csiro.pathling.fhirpath.collection.mixed;
 
 import au.csiro.pathling.fhirpath.collection.Collection;
-import au.csiro.pathling.fhirpath.column.NullRepresentation;
+import au.csiro.pathling.fhirpath.column.EmptyRepresentation;
 import au.csiro.pathling.fhirpath.definition.ChoiceChildDefinition;
 import java.util.Optional;
 import javax.annotation.Nonnull;
@@ -16,7 +16,7 @@
 public abstract class MixedCollection extends Collection {
 
   protected MixedCollection() {
-    super(NullRepresentation.getInstance(), Optional.empty(), Optional.empty(), Optional.empty());
+    super(EmptyRepresentation.getInstance(), Optional.empty(), Optional.empty(), Optional.empty());
   }
 
   /**

fhirpath/src/main/java/au/csiro/pathling/fhirpath/column/NullRepresentation.java → fhirpath/src/main/java/au/csiro/pathling/fhirpath/column/EmptyRepresentation.java

@@ -14,13 +14,13 @@
  * @author Piotr Szul
  * @author John Grimes
  */
-public class NullRepresentation extends ColumnRepresentation {
+public class EmptyRepresentation extends ColumnRepresentation {
 
-  static final ColumnRepresentation INSTANCE = new NullRepresentation();
+  static final ColumnRepresentation INSTANCE = new EmptyRepresentation();
   static final Column NULL_LITERAL = lit(null);
 
   /**
-   * @return a singleton instance of this class
+   * @return A singleton instance of this class
    */
   @Nonnull
   public static ColumnRepresentation getInstance() {
@@ -52,13 +52,13 @@ public ColumnRepresentation flatten() {
 
   @Nonnull
   @Override
-  public NullRepresentation traverse(@Nonnull final String fieldName) {
+  public EmptyRepresentation traverse(@Nonnull final String fieldName) {
     return this;
   }
 
   @Nonnull
   @Override
-  public NullRepresentation traverse(@Nonnull final String fieldName,
+  public EmptyRepresentation traverse(@Nonnull final String fieldName,
       final Optional<FHIRDefinedType> fhirType) {
     return traverse(fieldName);
   }