fix: Preserve extension map in ResourceCollection.copyWith()

piotrszul · claude · piotrszul · commit ec957f922384 · 2026-03-02T14:57:58.000+10:00
ResourceCollection.copyWith() was re-deriving the extension map from the
new column representation, which fails for non-resource structs (like
Extension) that lack the _extension field. This caused repeatAll(extension)
to only return top-level extensions, missing nested sub-extensions.

Added a private constructor that accepts an explicit extension map column,
and updated copyWith() to use it. Also consolidated RepeatAllFunctionDslTest
from 13 methods to 7 by chaining groups that share the same resource setup,
and refined the delta spec for same-type recursion depth scenarios.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
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
@@ -31,6 +31,7 @@
 import jakarta.annotation.Nonnull;
 import java.util.Optional;
 import lombok.Getter;
+import org.apache.spark.sql.Column;
 import org.hl7.fhir.exceptions.FHIRException;
 import org.hl7.fhir.r4.model.Enumerations.FHIRDefinedType;
 import org.hl7.fhir.r4.model.Enumerations.ResourceType;
@@ -47,7 +48,7 @@ public class ResourceCollection extends Collection {
   @Nonnull private final ResourceDefinition resourceDefinition;
 
   /**
-   * Creates a new ResourceCollection.
+   * Creates a new ResourceCollection, deriving the extension map from the column representation.
    *
    * @param columnRepresentation the column representation
    * @param type the FhirPath type
@@ -71,6 +72,29 @@ protected ResourceCollection(
     this.resourceDefinition = resourceDefinition;
   }
 
+  /**
+   * Creates a new ResourceCollection with an explicit extension map column. This is used by {@link
+   * #copyWith(ColumnRepresentation)} to preserve the resource-level extension map when creating
+   * copies with a different column representation (e.g. during repeatAll traversal).
+   *
+   * @param columnRepresentation the column representation
+   * @param type the FhirPath type
+   * @param fhirType the FHIR type
+   * @param definition the node definition
+   * @param resourceDefinition the resource definition
+   * @param extensionMapColumn the extension map column to preserve
+   */
+  private ResourceCollection(
+      @Nonnull final ColumnRepresentation columnRepresentation,
+      @Nonnull final Optional<FhirPathType> type,
+      @Nonnull final Optional<FHIRDefinedType> fhirType,
+      @Nonnull final Optional<? extends NodeDefinition> definition,
+      @Nonnull final ResourceDefinition resourceDefinition,
+      @Nonnull final Optional<Column> extensionMapColumn) {
+    super(columnRepresentation, type, fhirType, definition, extensionMapColumn);
+    this.resourceDefinition = resourceDefinition;
+  }
+
   @Nonnull
   private static Optional<FHIRDefinedType> getFhirType(@Nonnull final ResourceType resourceType) {
     return getFhirType(resourceType.toCode());
@@ -175,7 +199,12 @@ protected ColumnRepresentation getFid() {
   @Override
   public Collection copyWith(@Nonnull final ColumnRepresentation newValue) {
     return new ResourceCollection(
-        newValue, getType(), getFhirType(), getDefinition(), resourceDefinition);
+        newValue,
+        getType(),
+        getFhirType(),
+        getDefinition(),
+        resourceDefinition,
+        getExtensionMapColumn());
   }
 
   /**
diff --git a/fhirpath/src/test/java/au/csiro/pathling/fhirpath/dsl/RepeatAllFunctionDslTest.java b/fhirpath/src/test/java/au/csiro/pathling/fhirpath/dsl/RepeatAllFunctionDslTest.java
@@ -120,8 +120,47 @@ private static Questionnaire createEmptyQuestionnaire() {
     return questionnaire;
   }
 
+  /**
+   * Creates a Patient with all fields needed by the patient projection tests: two names, gender,
+   * marital status, and nested extensions.
+   *
+   * <p>Extension structure:
+   *
+   * <pre>
+   *   Patient
+   *     name: Smith, Jones
+   *     gender: MALE
+   *     maritalStatus: M (Married)
+   *     extension "http://example.com/simple" = "simple-value"
+   *     extension "http://example.com/parent"
+   *       extension "http://example.com/child" = "child-value"
+   * </pre>
+   *
+   * @return a Patient resource with all fields populated
+   */
+  private static Patient createPatient() {
+    final Patient patient = new Patient();
+    patient.setId("test-patient");
+    patient.addName().setFamily("Smith");
+    patient.addName().setFamily("Jones");
+    patient.setGender(AdministrativeGender.MALE);
+    patient.setMaritalStatus(
+        new CodeableConcept(new Coding("http://example.com/cs", "M", "Married")));
+
+    // Simple leaf extension.
+    patient.addExtension(
+        new Extension("http://example.com/simple", new StringType("simple-value")));
+
+    // Complex extension containing a nested sub-extension.
+    final Extension parent = new Extension("http://example.com/parent");
+    parent.addExtension(new Extension("http://example.com/child", new StringType("child-value")));
+    patient.addExtension(parent);
+
+    return patient;
+  }
+
   @FhirPathTest
-  public Stream<DynamicTest> testRepeatAllBasicTraversal() {
+  public Stream<DynamicTest> testRepeatAllQuestionnaireTraversal() {
     return builder()
         .withResource(createQuestionnaire())
         .group("repeatAll() basic traversal")
@@ -133,18 +172,54 @@ public Stream<DynamicTest> testRepeatAllBasicTraversal() {
             4,
             "repeatAll(item).count()",
             "repeatAll(item).count() returns total items across all levels")
+        .group("repeatAll() with filtering")
+        .testEquals(
+            List.of("1", "1.1"),
+            "repeatAll(item).where(type = 'group').linkId",
+            "repeatAll(item).where(type = 'group') filters items from all nesting levels")
+        .group("repeatAll() $index not defined")
+        .testError(
+            "repeatAll($index)", "$index is not available within repeatAll projection expression")
         .build();
   }
 
   @FhirPathTest
-  public Stream<DynamicTest> testRepeatAllFiltering() {
+  public Stream<DynamicTest> testRepeatAllPatientProjections() {
+    // repeatAll(first()) produces the same SQL DataType at every level (ArrayType(HumanName)),
+    // triggering same-type depth limiting. With a depth limit of 10, the initial level (where
+    // parentType is None and does not consume budget) plus 10 same-type levels produce 11 total
+    // copies of the first element before traversal stops.
     return builder()
-        .withResource(createQuestionnaire())
-        .group("repeatAll() with filtering")
+        .withResource(createPatient())
+        .group("repeatAll() non-recursive projection")
         .testEquals(
-            List.of("1", "1.1"),
-            "repeatAll(item).where(type = 'group').linkId",
-            "repeatAll(item).where(type = 'group') filters items from all nesting levels")
+            List.of("Smith", "Jones"),
+            "repeatAll(name).family",
+            "repeatAll(name).family returns same result as select() for non-recursive fields")
+        .group("repeatAll() singular primitive projection")
+        .testEquals(
+            "male",
+            "repeatAll(gender)",
+            "repeatAll(gender) returns a singular primitive value like select()")
+        .group("repeatAll() singular complex projection")
+        .testEquals(
+            "M",
+            "repeatAll(maritalStatus).coding.code",
+            "repeatAll(maritalStatus) returns a singular complex value with sub-elements intact")
+        .group("repeatAll() identity-like same-type depth limiting")
+        .testEquals(
+            11,
+            "name.repeatAll(first()).count()",
+            "repeatAll(first()) stops after same-type depth limit and returns bounded results")
+        .group("repeatAll() extension traversal")
+        .testEquals(
+            List.of(
+                "http://example.com/simple",
+                "http://example.com/parent",
+                "http://example.com/child"),
+            "repeatAll(extension).url",
+            "repeatAll(extension) recursively collects all extensions including nested"
+                + " sub-extensions")
         .build();
   }
 
@@ -157,6 +232,10 @@ public Stream<DynamicTest> testRepeatAllEmptyInput() {
             "repeatAll(item)", "repeatAll() on a resource with no items returns empty collection")
         .testEquals(
             0, "repeatAll(item).count()", "repeatAll(item).count() returns 0 for empty input")
+        .group("repeatAll() empty literal")
+        .testEmpty(
+            "{}.repeatAll($this)",
+            "repeatAll() on the FHIRPath empty literal returns empty collection")
         .build();
   }
 
@@ -231,56 +310,6 @@ public Stream<DynamicTest> testRepeatAllDuplicatesPreserved() {
         .build();
   }
 
-  @FhirPathTest
-  public Stream<DynamicTest> testRepeatAllNonRecursiveProjection() {
-    final Patient patient = new Patient();
-    patient.setId("test-patient");
-    patient.addName().setFamily("Smith");
-    patient.addName().setFamily("Jones");
-
-    return builder()
-        .withResource(patient)
-        .group("repeatAll() non-recursive projection")
-        .testEquals(
-            List.of("Smith", "Jones"),
-            "repeatAll(name).family",
-            "repeatAll(name).family returns same result as select() for non-recursive fields")
-        .build();
-  }
-
-  @FhirPathTest
-  public Stream<DynamicTest> testRepeatAllSingularPrimitive() {
-    final Patient patient = new Patient();
-    patient.setId("singular-patient");
-    patient.setGender(AdministrativeGender.MALE);
-
-    return builder()
-        .withResource(patient)
-        .group("repeatAll() singular primitive projection")
-        .testEquals(
-            "male",
-            "repeatAll(gender)",
-            "repeatAll(gender) returns a singular primitive value like select()")
-        .build();
-  }
-
-  @FhirPathTest
-  public Stream<DynamicTest> testRepeatAllSingularComplex() {
-    final Patient patient = new Patient();
-    patient.setId("complex-patient");
-    patient.setMaritalStatus(
-        new CodeableConcept(new Coding("http://example.com/cs", "M", "Married")));
-
-    return builder()
-        .withResource(patient)
-        .group("repeatAll() singular complex projection")
-        .testEquals(
-            "M",
-            "repeatAll(maritalStatus).coding.code",
-            "repeatAll(maritalStatus) returns a singular complex value with sub-elements intact")
-        .build();
-  }
-
   /**
    * Creates a QuestionnaireResponse with items nested through answer elements, to verify that
    * repeatAll can traverse through intermediate types (Item -> Answer -> Item).
diff --git a/openspec/changes/repeat-all-function/specs/fhirpath-repeat-all/spec.md b/openspec/changes/repeat-all-function/specs/fhirpath-repeat-all/spec.md
@@ -68,11 +68,21 @@ across all current items produces an empty collection (no new results).
 
 ### Requirement: repeatAll function handles empty input
 
-If the input collection is empty, the result SHALL be empty.
+If the input collection is empty, the result SHALL be empty. This SHALL hold
+both when the input is a typed collection that evaluates to empty at runtime
+(e.g., a resource where the traversed field is absent) and when the input is
+the FHIRPath empty literal `{}`.
 
-#### Scenario: Empty input collection
+#### Scenario: Empty collection from absent field
 
-- **WHEN** `repeatAll` is called on an empty collection
+- **WHEN** `repeatAll` is called on a resource where the traversed field is
+  absent (e.g., a Questionnaire with no items)
+- **THEN** the result SHALL be an empty collection
+
+#### Scenario: Empty collection literal
+
+- **WHEN** `repeatAll` is called on the FHIRPath empty literal `{}` (e.g.,
+  `{}.repeatAll(item)`)
 - **THEN** the result SHALL be an empty collection
 
 ### Requirement: repeatAll function supports chained expressions
@@ -104,26 +114,53 @@ subsequent FHIRPath operations such as path navigation, `where()`, `count()`,
 
 ### Requirement: repeatAll function bounds same-type recursion depth
 
-The function SHALL always navigate to the deepest level possible regardless of
-the encoding's maximum nesting level. Recursion depth SHALL only be limited when
-the traversal expression is self-referential (i.e. navigates to a node of the
-same type), in which case a hardcoded depth limit of 10 SHALL prevent infinite
-loops. This matches the approach used by the SQL on FHIR `repeat` clause
-implementation.
-
-#### Scenario: Same-type recursion bounded by depth limit
-
-- **WHEN** `repeatAll` is evaluated and the traversal expression navigates to a
-  node of the same type (e.g., `item` navigating from Item to Item)
-- **THEN** the function SHALL stop same-type traversal after the hardcoded depth
-  limit and return results collected up to that point
-
-#### Scenario: Cross-type traversal not limited by depth
-
-- **WHEN** `repeatAll` is evaluated and the traversal expression crosses
-  different types before recurring (e.g., Item → Answer → Item)
-- **THEN** the cross-type steps SHALL NOT consume depth budget, allowing deeper
-  traversal than the same-type limit alone would permit
+The function delegates recursive traversal to `UnresolvedTransformTree`, which
+limits recursion depth based on Spark SQL `DataType` equality. The depth counter
+only decrements when the traversal expression produces a node whose resolved SQL
+`DataType` is structurally identical to its parent's. When the counter reaches
+zero, traversal stops and returns results collected up to that point. The
+hardcoded same-type depth limit is 10, matching the default
+`maxUnboundTraversalDepth` used by the SQL on FHIR `repeat` clause.
+
+For traversals through schema-truncated structures (e.g.,
+`Questionnaire.repeatAll(item)`), each nesting level has a progressively smaller
+`StructType` due to the encoding's `maxNestingLevel` truncation. These are NOT
+same-type traversals — the depth counter is never decremented, and traversal
+terminates naturally when a field is no longer present in the truncated schema
+(`FIELD_NOT_FOUND`).
+
+Same-type recursion occurs when the traversal produces an element with an
+identical SQL `DataType` at every level. Known cases include:
+
+- **Extension traversal**: Extension structs have a fixed flat schema at all
+  nesting levels (extensions use a global `_extension` map rather than recursive
+  struct fields), so `repeatAll(extension)` produces same-type nodes.
+- **Identity-like traversals**: Functions such as `first()` that return an
+  element of the same type as the input, causing the SQL `DataType` to remain
+  unchanged across levels.
+
+#### Scenario: Extension traversal bounded by actual depth or limit
+
+- **WHEN** `repeatAll(extension)` is evaluated on a resource whose elements
+  carry nested extensions
+- **THEN** the function SHALL traverse extensions up to the actual nesting depth
+  or the same-type depth limit (10), whichever is reached first, and return all
+  collected extensions
+
+#### Scenario: Identity-like traversal bounded by depth limit
+
+- **WHEN** `repeatAll` is evaluated with a traversal that always produces the
+  same SQL `DataType` (e.g., `repeatAll(first())` on a complex-typed
+  collection)
+- **THEN** the function SHALL stop after the same-type depth limit and return
+  results collected up to that point
+
+#### Scenario: Schema-truncated traversal terminates naturally
+
+- **WHEN** `repeatAll(item)` is evaluated on a Questionnaire with nested items
+- **THEN** traversal SHALL terminate naturally when the schema no longer
+  contains the `item` field (due to `maxNestingLevel` truncation during the
+  encoding), without consuming same-type depth budget
 
 ### Requirement: repeatAll function does not define $index
 
