Simplified the grammar for FILLNULL

Signed-off-by: Norman Jordan <[email protected]>
Bit-Quill · Dec 4, 2024 · 8559e1c · 8559e1c
1 parent 0f04940
commit 8559e1c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 13 deletions.
diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4
@@ -134,19 +134,15 @@ fillnullCommand
    ;
 
 fillNullWithTheSameValue
-   : WITH nullReplacement IN nullableField (COMMA nullableField)*
+   : WITH nullReplacement = valueExpression IN nullableFieldList = fieldList
    ;
 
 fillNullWithFieldVariousValues
-   : USING nullableField EQUAL nullReplacement (COMMA nullableField EQUAL nullReplacement)*
+   : USING nullReplacementExpression (COMMA nullReplacementExpression)*
    ;
 
-nullableField
-   : fieldExpression
-   ;
-
-nullReplacement
-   : valueExpression
+nullReplacementExpression
+   : nullableField = fieldExpression EQUAL nullReplacement = valueExpression
    ;
 
 kmeansCommand

diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java
@@ -399,20 +399,22 @@ public UnresolvedPlan visitFillNullWithTheSameValue(
       OpenSearchPPLParser.FillNullWithTheSameValueContext ctx) {
     return new FillNull(
         FillNull.ContainNullableFieldFill.ofSameValue(
-            internalVisitExpression(ctx.nullReplacement()),
-            ctx.nullableField().stream().map(f -> (Field) internalVisitExpression(f)).toList()));
+            internalVisitExpression(ctx.nullReplacement),
+            ctx.nullableFieldList.fieldExpression().stream()
+                .map(f -> (Field) internalVisitExpression(f))
+                .toList()));
   }
 
   /** fillnull command. */
   @Override
   public UnresolvedPlan visitFillNullWithFieldVariousValues(
       OpenSearchPPLParser.FillNullWithFieldVariousValuesContext ctx) {
     ImmutableList.Builder<FillNull.NullableFieldFill> replacementsBuilder = ImmutableList.builder();
-    for (int i = 0; i < ctx.nullableField().size(); i++) {
+    for (int i = 0; i < ctx.nullReplacementExpression().size(); i++) {
       replacementsBuilder.add(
           new FillNull.NullableFieldFill(
-              (Field) internalVisitExpression(ctx.nullableField(i)),
-              internalVisitExpression(ctx.nullReplacement(i))));
+              (Field) internalVisitExpression(ctx.nullReplacementExpression(i).nullableField),
+              internalVisitExpression(ctx.nullReplacementExpression(i).nullReplacement)));
     }
 
     return new FillNull(