add parsing (only) of container-phase in schema

3 files changed, 84 insertions, 2 deletions

diff --git a/config-model/src/main/java/com/yahoo/schema/parser/ParsedRankProfile.java b/config-model/src/main/java/com/yahoo/schema/parser/ParsedRankProfile.java
index 64dd8dd0ad4..2de7fab09f4 100644
--- a/config-model/src/main/java/com/yahoo/schema/parser/ParsedRankProfile.java
+++ b/config-model/src/main/java/com/yahoo/schema/parser/ParsedRankProfile.java
@@ -53,6 +53,8 @@ class ParsedRankProfile extends ParsedBlock {
     private final Map<Reference, RankProfile.Constant> constants = new LinkedHashMap<>();
     private final Map<Reference, RankProfile.Input> inputs = new LinkedHashMap<>();
     private final List<OnnxModel> onnxModels = new ArrayList<>();
+    private Integer containerPhaseRerankCount = null;
+    private String containerPhaseExpression = null;
 
     ParsedRankProfile(String name) {
         super(name, "rank-profile");
@@ -77,6 +79,8 @@ class ParsedRankProfile extends ParsedBlock {
     List<ParsedRankFunction> getFunctions() { return List.copyOf(functions.values()); }
     List<MutateOperation> getMutateOperations() { return List.copyOf(mutateOperations); }
     List<String> getInherited() { return List.copyOf(inherited); }
+    Optional<Integer> getContainerPhaseRerankCount() { return Optional.ofNullable(this.containerPhaseRerankCount); }
+    Optional<String> getContainerPhaseExpression() { return Optional.ofNullable(this.containerPhaseExpression); }
 
     Map<String, Boolean> getFieldsWithRankFilter() { return Collections.unmodifiableMap(fieldsRankFilter); }
     Map<String, Integer> getFieldsWithRankWeight() { return Collections.unmodifiableMap(fieldsRankWeight); }
@@ -197,11 +201,21 @@ class ParsedRankProfile extends ParsedBlock {
         this.secondPhaseExpression = expression;
     }
 
+    void setContainerPhaseExpression(String expression) {
+        verifyThat(containerPhaseExpression == null, "already has container-phase expression");
+        this.containerPhaseExpression = expression;
+    }
+
+    void setContainerPhaseRerankCount(int count) {
+        verifyThat(containerPhaseRerankCount == null, "already has container-phase rerank-count");
+        this.containerPhaseRerankCount = count;
+    }
+
     void setStrict(boolean strict) {
         verifyThat(this.strict == null, "already has strict");
         this.strict = strict;
     }
-  
+
     void setTermwiseLimit(double limit) {
         verifyThat(termwiseLimit == null, "already has termwise-limit");
         this.termwiseLimit = limit;
diff --git a/config-model/src/main/javacc/SchemaParser.jj b/config-model/src/main/javacc/SchemaParser.jj
index a9c67a0bb60..888e5ffb665 100644
--- a/config-model/src/main/javacc/SchemaParser.jj
+++ b/config-model/src/main/javacc/SchemaParser.jj
@@ -279,6 +279,7 @@ TOKEN :
 | < MAXHITS: "max-hits" >
 | < FIRSTPHASE: "first-phase" >
 | < SECONDPHASE: "second-phase" >
+| < CONTAINERPHASE: "container-phase" >
 | < MACRO: "macro" >
 | < INLINE: "inline" >
 | < ARITY: "arity" >
@@ -1706,6 +1707,7 @@ void rankProfileItem(ParsedSchema schema, ParsedRankProfile profile) : { }
       | rankFeatures(profile)
       | rankProperties(profile)
       | secondPhase(profile)
+      | containerPhase(profile)
       | inputs(profile)
       | constants(schema, profile)
       | matchFeatures(profile)
@@ -1923,6 +1925,34 @@ void secondPhaseItem(ParsedRankProfile profile) :
     )
 }
 
+/**
+ * Consumes the container-phase block of a rank profile.
+ *
+ * @param profile The rank profile to modify.
+ */
+void containerPhase(ParsedRankProfile profile) : { }
+{
+    <CONTAINERPHASE> lbrace() (containerPhaseItem(profile) (<NL>)*)* <RBRACE>
+}
+
+/**
+ * Consumes a statement for a container-phase block.
+ *
+ * @param profile The rank profile to modify.
+ */
+void containerPhaseItem(ParsedRankProfile profile) :
+{
+    String expression;
+    int rerankCount;
+}
+{
+    ( expression = expression()                       { profile.setContainerPhaseExpression(expression); }
+    | (<RERANKCOUNT> <COLON> rerankCount = integer()) { profile.setContainerPhaseRerankCount(rerankCount); }
+    )
+}
+
+
+
 /** Consumes an inputs block of a rank profile. */
 void inputs(ParsedRankProfile profile) :
 {
@@ -2519,7 +2549,7 @@ String expression() :
     ( <EXPRESSION_SL> { exp = token.image.substring(token.image.indexOf(":") + 1); } |
       <EXPRESSION_ML> { exp = token.image.substring(token.image.indexOf("{") + 1,
                                                     token.image.lastIndexOf("}")); } )
-    { return exp; }
+    { return exp.trim(); }
 }
 
 String identifierWithDash() :
@@ -2555,6 +2585,7 @@ String identifier() : { }
       | <CONSTANT>
       | <CONSTANTS>
       | <CONTEXT>
+      | <CONTAINERPHASE>
       | <CREATEIFNONEXISTENT>
       | <DENSEPOSTINGLISTTHRESHOLD>
       | <DESCENDING>
diff --git a/config-model/src/test/java/com/yahoo/schema/parser/SchemaParserTestCase.java b/config-model/src/test/java/com/yahoo/schema/parser/SchemaParserTestCase.java
index 150c237bbba..d08c2266370 100644
--- a/config-model/src/test/java/com/yahoo/schema/parser/SchemaParserTestCase.java
+++ b/config-model/src/test/java/com/yahoo/schema/parser/SchemaParserTestCase.java
@@ -83,6 +83,43 @@ public class SchemaParserTestCase {
                                 "}\n")).getMessage());
     }
 
+    @Test
+    void container_phase_can_be_parsed() throws Exception {
+        String input = """
+            schema foo {
+                rank-profile normal {
+                    first-phase {
+                        expression {
+                            rankingExpression(1.0)
+                        }
+                    }
+                }
+                rank-profile bar {
+                    container-phase {
+                        expression: onnx(mymodel)
+                        rerank-count: 79
+                    }
+                }
+            }
+            """;
+        ParsedSchema schema = parseString(input);
+        assertEquals("foo", schema.name());
+        var rplist = schema.getRankProfiles();
+        assertEquals(2, rplist.size());
+        var rp0 = rplist.get(0);
+        assertEquals("normal", rp0.name());
+        assertFalse(rp0.getContainerPhaseRerankCount().isPresent());
+        assertFalse(rp0.getContainerPhaseExpression().isPresent());
+        assertTrue(rp0.getFirstPhaseExpression().isPresent());
+        assertEquals("rankingExpression(1.0)", rp0.getFirstPhaseExpression().get());
+        var rp1 = rplist.get(1);
+        assertEquals("bar", rp1.name());
+        assertTrue(rp1.getContainerPhaseRerankCount().isPresent());
+        assertTrue(rp1.getContainerPhaseExpression().isPresent());
+        assertEquals(79, rp1.getContainerPhaseRerankCount().get());
+        assertEquals("onnx(mymodel)", rp1.getContainerPhaseExpression().get());
+    }
+
     void checkFileParses(String fileName) throws Exception {
         var schema = parseFile(fileName);
         assertNotNull(schema);