Updated grammar for FILLNULL to match grammar in SQL project (opensea…

…rch-project#968) Signed-off-by: Norman Jordan <[email protected]>
Bit-Quill · Dec 11, 2024 · 64035d0 · 64035d0
1 parent 876bd04
commit 64035d0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 18 deletions.
diff --git a/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 b/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4
@@ -233,21 +233,16 @@ fillnullCommand
    | fillNullWithFieldVariousValues)
    ;
 
- fillNullWithTheSameValue
- : WITH nullReplacement IN nullableField (COMMA nullableField)*
- ;
-
- fillNullWithFieldVariousValues
- : USING nullableField EQUAL nullReplacement (COMMA nullableField EQUAL nullReplacement)*
- ;
-
+fillNullWithTheSameValue
+   : WITH nullReplacement = valueExpression IN nullableFieldList = fieldList
+   ;
 
-   nullableField
-   : fieldExpression
+fillNullWithFieldVariousValues
+   : USING nullableReplacementExpression (COMMA nullableReplacementExpression)*
    ;
 
-   nullReplacement
-   : expression
+nullableReplacementExpression
+   : nullableField = fieldExpression EQUAL nullableReplacement = valueExpression
    ;
 
 flattenCommand

diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java
@@ -589,19 +589,18 @@ public UnresolvedPlan visitFillnullCommand(OpenSearchPPLParser.FillnullCommandCo
     FillNullWithFieldVariousValuesContext variousValuesContext = ctx.fillNullWithFieldVariousValues();
     if (sameValueContext != null) {
       // todo consider using expression instead of Literal
-      UnresolvedExpression replaceNullWithMe = internalVisitExpression(sameValueContext.nullReplacement().expression());
-      List<Field> fieldsToReplace = sameValueContext.nullableField()
+      UnresolvedExpression replaceNullWithMe = internalVisitExpression(sameValueContext.nullReplacement);
+      List<Field> fieldsToReplace = sameValueContext.nullableFieldList.fieldExpression()
               .stream()
               .map(this::internalVisitExpression)
               .map(Field.class::cast)
               .collect(Collectors.toList());
       return new FillNull(ofSameValue(replaceNullWithMe, fieldsToReplace));
     } else if (variousValuesContext != null) {
-      List<NullableFieldFill> nullableFieldFills = IntStream.range(0, variousValuesContext.nullableField().size())
+      List<NullableFieldFill> nullableFieldFills = IntStream.range(0, variousValuesContext.nullableReplacementExpression().size())
               .mapToObj(index -> {
-                variousValuesContext.nullableField(index);
-                UnresolvedExpression replaceNullWithMe = internalVisitExpression(variousValuesContext.nullReplacement(index).expression());
-                Field nullableFieldReference = (Field) internalVisitExpression(variousValuesContext.nullableField(index));
+                UnresolvedExpression replaceNullWithMe = internalVisitExpression(variousValuesContext.nullableReplacementExpression(index).nullableReplacement);
+                Field nullableFieldReference = (Field) internalVisitExpression(variousValuesContext.nullableReplacementExpression(index).nullableField);
                 return new NullableFieldFill(nullableFieldReference, replaceNullWithMe);
               })
               .collect(Collectors.toList());