Add rename PPL function

integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLRenameITSuite.scala

@@ -0,0 +1,152 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.flint.spark.ppl
+
+import org.opensearch.sql.ppl.utils.DataTypeTransformer.seq
+
+import org.apache.spark.sql.{QueryTest, Row}
+import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar}
+import org.apache.spark.sql.catalyst.expressions.Alias
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Project}
+import org.apache.spark.sql.streaming.StreamTest
+
+class FlintSparkPPLRenameITSuite
+    extends QueryTest
+    with LogicalPlanTestUtils
+    with FlintPPLSuite
+    with StreamTest {
+
+  /** Test table and index name */
+  private val testTable = "spark_catalog.default.flint_ppl_test"
+
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+
+    // Create test table
+    createPartitionedStateCountryTable(testTable)
+  }
+
+  protected override def afterEach(): Unit = {
+    super.afterEach()
+    // Stop all streaming jobs if any
+    spark.streams.active.foreach { job =>
+      job.stop()
+      job.awaitTermination()
+    }
+  }
+
+  test("test single renamed field in fields command") {
+    val frame = sql(s"""
+         | source = $testTable | rename age as renamed_age | fields name, renamed_age
+         | """.stripMargin)
+
+    // Retrieve the results
+    val results: Array[Row] = frame.collect()
+    // Define the expected results
+    val expectedResults: Array[Row] =
+      Array(Row("Jake", 70), Row("Hello", 30), Row("John", 25), Row("Jane", 20))
+    // Compare the results
+    implicit val rowOrdering: Ordering[Row] = Ordering.by[Row, String](_.getAs[String](0))
+    assert(results.sorted.sameElements(expectedResults.sorted))
+
+    // Retrieve the logical plan
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    // Define the expected logical plan
+    val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
+    val fieldsProjectList = Seq(UnresolvedAttribute("name"), UnresolvedAttribute("renamed_age"))
+    val renameProjectList =
+      Seq(UnresolvedStar(None), Alias(UnresolvedAttribute("age"), "renamed_age")())
+    val expectedPlan = Project(fieldsProjectList, Project(renameProjectList, table))
+    // Compare the two plans
+    comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
+  }
+
+  test("test multiple renamed fields in fields command") {
+    val frame = sql(s"""
+         | source = $testTable | rename name as renamed_name, country as renamed_country | fields renamed_name, age, renamed_country
+         | """.stripMargin)
+
+    val results: Array[Row] = frame.collect()
+    // results.foreach(println(_))
+    val expectedResults: Array[Row] =
+      Array(
+        Row("Jake", 70, "USA"),
+        Row("Hello", 30, "USA"),
+        Row("John", 25, "Canada"),
+        Row("Jane", 20, "Canada"))
+    implicit val rowOrdering: Ordering[Row] = Ordering.by[Row, String](_.getAs[String](0))
+    assert(results.sorted.sameElements(expectedResults.sorted))
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
+    val fieldsProjectList = Seq(
+      UnresolvedAttribute("renamed_name"),
+      UnresolvedAttribute("age"),
+      UnresolvedAttribute("renamed_country"))
+    val renameProjectList =
+      Seq(
+        UnresolvedStar(None),
+        Alias(UnresolvedAttribute("name"), "renamed_name")(),
+        Alias(UnresolvedAttribute("country"), "renamed_country")())
+    val expectedPlan = Project(fieldsProjectList, Project(renameProjectList, table))
+    comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
+  }
+
+  test("test renamed fields without fields command") {
+    val frame = sql(s"""
+         | source = $testTable | rename age as user_age, country as user_country
+         | """.stripMargin)
+
+    val results: Array[Row] = frame.collect()
+    val expectedResults: Array[Row] = Array(
+      Row("Jake", 70, "California", "USA", 2023, 4, 70, "USA"),
+      Row("Hello", 30, "New York", "USA", 2023, 4, 30, "USA"),
+      Row("John", 25, "Ontario", "Canada", 2023, 4, 25, "Canada"),
+      Row("Jane", 20, "Quebec", "Canada", 2023, 4, 20, "Canada"))
+    implicit val rowOrdering: Ordering[Row] = Ordering.by[Row, String](_.getAs[String](0))
+    assert(results.sorted.sameElements(expectedResults.sorted))
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
+    val renameProjectList = Seq(
+      UnresolvedStar(None),
+      Alias(UnresolvedAttribute("age"), "user_age")(),
+      Alias(UnresolvedAttribute("country"), "user_country")())
+    val expectedPlan = Project(seq(UnresolvedStar(None)), Project(renameProjectList, table))
+    comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
+  }
+
+  test("test renamed field used in aggregation") {
+    val frame = sql(s"""
+         | source = $testTable | rename age as user_age | stats avg(user_age) by country
+         | """.stripMargin)
+
+    val results: Array[Row] = frame.collect()
+    val expectedResults: Array[Row] = Array(Row(22.5, "Canada"), Row(50.0, "USA"))
+    implicit val rowOrdering: Ordering[Row] = Ordering.by[Row, Double](_.getAs[Double](0))
+    assert(results.sorted.sameElements(expectedResults.sorted))
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
+    val renameProjectList =
+      Seq(UnresolvedStar(None), Alias(UnresolvedAttribute("age"), "user_age")())
+    val aggregateExpressions =
+      Seq(
+        Alias(
+          UnresolvedFunction(
+            Seq("AVG"),
+            Seq(UnresolvedAttribute("user_age")),
+            isDistinct = false),
+          "avg(user_age)")(),
+        Alias(UnresolvedAttribute("country"), "country")())
+    val aggregatePlan = Aggregate(
+      Seq(Alias(UnresolvedAttribute("country"), "country")()),
+      aggregateExpressions,
+      Project(renameProjectList, table))
+    val expectedPlan = Project(seq(UnresolvedStar(None)), aggregatePlan)
+    comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
+  }
+}

ppl-spark-integration/README.md

@@ -404,6 +404,10 @@ Limitation: Overriding existing field is unsupported, following queries throw ex
 - `source=apache | patterns new_field='no_numbers' pattern='[0-9]' message | fields message, no_numbers`
 - `source=apache | patterns new_field='no_numbers' pattern='[0-9]' message | stats count() by no_numbers`
 
+**Rename**
+- `source=accounts | rename email as user_email | fields id, user_email`
+- `source=accounts | rename id as user_id, email as user_email | fields user_id, user_email`
+
 _- **Limitation: Overriding existing field is unsupported:**_
   - `source=accounts | grok address '%{NUMBER} %{GREEDYDATA:address}' | fields address`
 

ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4

@@ -46,6 +46,7 @@ commands
    | parseCommand
    | patternsCommand
    | lookupCommand
+   | renameCommand
    ;
 
 searchCommand

ppl-spark-integration/src/main/java/org/opensearch/sql/ast/tree/Rename.java

@@ -11,7 +11,7 @@
 import lombok.RequiredArgsConstructor;
 import lombok.ToString;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
-import org.opensearch.sql.ast.expression.Map;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
 
 import java.util.List;
 
@@ -20,10 +20,10 @@
 @Getter
 @RequiredArgsConstructor
 public class Rename extends UnresolvedPlan {
-  private final List<Map> renameList;
+  private final List<UnresolvedExpression> renameList;
   private UnresolvedPlan child;
 
-  public Rename(List<Map> renameList, UnresolvedPlan child) {
+  public Rename(List<UnresolvedExpression> renameList, UnresolvedPlan child) {
     this.renameList = renameList;
     this.child = child;
   }

ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java

@@ -72,6 +72,7 @@
 import org.opensearch.sql.ast.tree.RareAggregation;
 import org.opensearch.sql.ast.tree.RareTopN;
 import org.opensearch.sql.ast.tree.Relation;
+import org.opensearch.sql.ast.tree.Rename;
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.SubqueryAlias;
 import org.opensearch.sql.ast.tree.TopAggregation;
@@ -399,6 +400,19 @@ public LogicalPlan visitParse(Parse node, CatalystPlanContext context) {
         return ParseStrategy.visitParseCommand(node, sourceField, parseMethod, arguments, pattern, context);
     }
 
+    @Override
+    public LogicalPlan visitRename(Rename node, CatalystPlanContext context) {
+        node.getChild().get(0).accept(this, context);
+        if (context.getNamedParseExpressions().isEmpty()) {
+            // Create an UnresolvedStar for all-fields projection
+            context.getNamedParseExpressions().push(UnresolvedStar$.MODULE$.apply(Option.empty()));
+        }
+        visitExpressionList(node.getRenameList(), context);
+        Seq<NamedExpression> projectExpressions = context.retainAllNamedParseExpressions(p -> (NamedExpression) p);
+        // build the plan with the projection step
+        return context.apply(p -> new org.apache.spark.sql.catalyst.plans.logical.Project(projectExpressions, p));
+    }
+
     @Override
     public LogicalPlan visitEval(Eval node, CatalystPlanContext context) {
         LogicalPlan child = node.getChild().get(0).accept(this, context);

ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java

@@ -22,7 +22,6 @@
 import org.opensearch.sql.ast.expression.FieldsMapping;
 import org.opensearch.sql.ast.expression.Let;
 import org.opensearch.sql.ast.expression.Literal;
-import org.opensearch.sql.ast.expression.Map;
 import org.opensearch.sql.ast.expression.ParseMethod;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.ast.expression.Scope;
@@ -211,9 +210,9 @@ public UnresolvedPlan visitRenameCommand(OpenSearchPPLParser.RenameCommandContex
         ctx.renameClasue().stream()
             .map(
                 ct ->
-                    new Map(
-                        internalVisitExpression(ct.orignalField),
-                        internalVisitExpression(ct.renamedField)))
+                    new Alias(
+                        ct.renamedField.getText(),
+                        internalVisitExpression(ct.orignalField)))
             .collect(Collectors.toList()));
   }
 

ppl-spark-integration/src/test/scala/org/opensearch/flint/spark/ppl/PPLLogicalPlanRenameTranslatorTestSuite.scala

@@ -0,0 +1,116 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.flint.spark.ppl
+
+import org.opensearch.flint.spark.ppl.PlaneUtils.plan
+import org.opensearch.sql.ppl.{CatalystPlanContext, CatalystQueryPlanVisitor}
+import org.opensearch.sql.ppl.utils.DataTypeTransformer.seq
+import org.scalatest.matchers.should.Matchers
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Descending, ExprId, NamedExpression, SortOrder}
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical.{Project, Sort}
+
+class PPLLogicalPlanRenameTranslatorTestSuite
+    extends SparkFunSuite
+    with PlanTest
+    with LogicalPlanTestUtils
+    with Matchers {
+
+  private val planTransformer = new CatalystQueryPlanVisitor()
+  private val pplParser = new PPLSyntaxParser()
+
+  test("test renamed fields not included in fields expressions") {
+    val context = new CatalystPlanContext
+    val logPlan =
+      planTransformer.visit(
+        plan(pplParser, "source=t | rename a as r_a, b as r_b | fields c"),
+        context)
+    val renameProjectList: Seq[NamedExpression] =
+      Seq(
+        UnresolvedStar(None),
+        Alias(UnresolvedAttribute("a"), "r_a")(),
+        Alias(UnresolvedAttribute("b"), "r_b")())
+    val expectedPlan = Project(
+      seq(UnresolvedAttribute("c")),
+      Project(renameProjectList, UnresolvedRelation(Seq("t"))))
+    comparePlans(expectedPlan, logPlan, checkAnalysis = false)
+  }
+
+  test("test renamed fields included in fields expression") {
+    val context = new CatalystPlanContext
+    val logPlan =
+      planTransformer.visit(
+        plan(pplParser, "source=t | rename a as r_a, b as r_b | fields r_a, r_b, c"),
+        context)
+
+    val renameProjectList: Seq[NamedExpression] =
+      Seq(
+        UnresolvedStar(None),
+        Alias(UnresolvedAttribute("a"), "r_a")(),
+        Alias(UnresolvedAttribute("b"), "r_b")())
+    val expectedPlan = Project(
+      seq(UnresolvedAttribute("r_a"), UnresolvedAttribute("r_b"), UnresolvedAttribute("c")),
+      Project(renameProjectList, UnresolvedRelation(Seq("t"))))
+    comparePlans(expectedPlan, logPlan, checkAnalysis = false)
+  }
+
+  test("test renamed fields without fields command") {
+    val context = new CatalystPlanContext
+    val logPlan =
+      planTransformer.visit(plan(pplParser, "source=t | rename a as r_a, b as r_b"), context)
+
+    val renameProjectList: Seq[NamedExpression] =
+      Seq(
+        UnresolvedStar(None),
+        Alias(UnresolvedAttribute("a"), "r_a")(),
+        Alias(UnresolvedAttribute("b"), "r_b")())
+    val expectedPlan =
+      Project(seq(UnresolvedStar(None)), Project(renameProjectList, UnresolvedRelation(Seq("t"))))
+    comparePlans(expectedPlan, logPlan, checkAnalysis = false)
+  }
+
+  test("test renamed fields with sort") {
+    val context = new CatalystPlanContext
+    val logPlan =
+      planTransformer.visit(
+        plan(pplParser, "source=t | rename a as r_a, b as r_b | sort - r_a | fields r_b"),
+        context)
+
+    val renameProjectList: Seq[NamedExpression] =
+      Seq(
+        UnresolvedStar(None),
+        Alias(UnresolvedAttribute("a"), "r_a")(),
+        Alias(UnresolvedAttribute("b"), "r_b")())
+    val renameProject = Project(renameProjectList, UnresolvedRelation(Seq("t")))
+    val sortOrder = SortOrder(UnresolvedAttribute("r_a"), Descending, Seq.empty)
+    val sort = Sort(seq(sortOrder), global = true, renameProject)
+    val expectedPlan = Project(seq(UnresolvedAttribute("r_b")), sort)
+    comparePlans(expectedPlan, logPlan, checkAnalysis = false)
+  }
+
+  test("test rename eval expression output") {
+    val context = new CatalystPlanContext
+    val logPlan =
+      planTransformer.visit(
+        plan(pplParser, "source=t | eval a = RAND() | rename a as eval_rand | fields eval_rand"),
+        context)
+
+    val evalProjectList: Seq[NamedExpression] = Seq(
+      UnresolvedStar(None),
+      Alias(UnresolvedFunction("rand", Seq.empty, isDistinct = false), "a")(
+        exprId = ExprId(0),
+        qualifier = Seq.empty))
+    val evalProject = Project(evalProjectList, UnresolvedRelation(Seq("t")))
+    val renameProjectList: Seq[NamedExpression] =
+      Seq(UnresolvedStar(None), Alias(UnresolvedAttribute("a"), "eval_rand")())
+    val expectedPlan =
+      Project(seq(UnresolvedAttribute("eval_rand")), Project(renameProjectList, evalProject))
+    comparePlans(expectedPlan, logPlan, checkAnalysis = false)
+  }
+}