From 64035d0f44123969327bfab154cb45aff9b050eb Mon Sep 17 00:00:00 2001
From: normanj-bitquill <78755797+normanj-bitquill@users.noreply.github.com>
Date: Thu, 5 Dec 2024 16:50:49 -0800
Subject: [PATCH] Updated grammar for FILLNULL to match grammar in SQL project
 (#968)

Signed-off-by: Norman Jordan <norman.jordan@improving.com>
---
 .../src/main/antlr4/OpenSearchPPLParser.g4    | 19 +++++++------------
 .../opensearch/sql/ppl/parser/AstBuilder.java | 11 +++++------
 2 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 b/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4
index 073e6a332..7a34f5840 100644
--- 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
index 1709b7cd7..2ea23babf 100644
--- 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());