fix(spark): enable aliased expressions to round-trip (#348)

A number of the TPC-DS tests were failing because the query
contains multiple aliases to the same expression, causing a
potential mismatch in the reference index.  Although the plans
were equivalent, the substrait POJO comparison failed.
This commit uses the `hint` field of the Rel message to store the
alias names, and restore them back to the Spark plan to match
the original.

Author: andrew-coleman

core/src/main/java/io/substrait/relation/ProtoRelConverter.java

@@ -423,7 +423,8 @@ protected Project newProject(ProjectRel rel) {
 
     builder
         .commonExtension(optionalAdvancedExtension(rel.getCommon()))
-        .remap(optionalRelmap(rel.getCommon()));
+        .remap(optionalRelmap(rel.getCommon()))
+        .hint(optionalHint(rel.getCommon()));
     if (rel.hasAdvancedExtension()) {
       builder.extension(advancedExtension(rel.getAdvancedExtension()));
     }

spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala

@@ -16,7 +16,7 @@
  */
 package io.substrait.spark.logical
 
-import io.substrait.spark.{DefaultRelVisitor, SparkExtension, ToSparkType}
+import io.substrait.spark.{DefaultRelVisitor, SparkExtension, ToSparkType, ToSubstraitType}
 import io.substrait.spark.expression._
 
 import org.apache.spark.sql.SparkSession
@@ -33,7 +33,7 @@ import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat
 import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, IntegerType, MapType, StructField, StructType}
+import org.apache.spark.sql.types.{DataType, IntegerType, StructField, StructType}
 
 import io.substrait.`type`.{NamedStruct, StringTypeVisitor, Type}
 import io.substrait.{expression => exp}
@@ -244,18 +244,43 @@ class ToLogicalPlan(spark: SparkSession) extends DefaultRelVisitor[LogicalPlan]
     }
   }
 
+  /**
+   * Returns the top level field (column) names for the given relation, if they have been specified
+   * in the optional `hint` message. Does not include the field names of any inner structs.
+   * @param rel
+   * @return
+   *   Optional list of names.
+   */
+  private def fieldNames(rel: relation.Rel): Option[Seq[String]] = {
+    if (rel.getHint.isPresent && !rel.getHint.get().getOutputNames.isEmpty) {
+      Some(
+        ToSubstraitType
+          .toNamedStruct(ToSparkType.toStructType(
+            NamedStruct.of(rel.getHint.get.getOutputNames, rel.getRecordType)))
+          .names
+          .asScala)
+    } else {
+      None
+    }
+  }
+
   override def visit(project: relation.Project): LogicalPlan = {
     val child = project.getInput.accept(this)
     val (output, createProject) = child match {
       case a: Aggregate => (a.aggregateExpressions, false)
       case other => (other.output, true)
     }
+    val names = fieldNames(project).getOrElse(List.empty)
 
     withOutput(output) {
-      val projectList =
+      val projectExprs =
         project.getExpressions.asScala
           .map(expr => expr.accept(expressionConverter))
-          .map(toNamedExpression)
+      val projectList = if (names.size == projectExprs.size) {
+        projectExprs.zip(names).map { case (expr, name) => Alias(expr, name)() }
+      } else {
+        projectExprs.map(toNamedExpression)
+      }
       if (createProject) {
         Project(projectList, child)
       } else {
@@ -267,7 +292,9 @@ class ToLogicalPlan(spark: SparkSession) extends DefaultRelVisitor[LogicalPlan]
 
   override def visit(expand: relation.Expand): LogicalPlan = {
     val child = expand.getInput.accept(this)
-    val names = expand.getHint.get().getOutputNames.asScala
+    val names = fieldNames(expand).getOrElse(
+      expand.getFields.asScala.zipWithIndex.map { case (_, i) => s"col$i" }
+    )
 
     withChild(child) {
       val projections = expand.getFields.asScala

spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala

@@ -173,10 +173,12 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging {
         }
     }
     val projects = projectExpressions.map(toExpression(newOutput))
+    val names = ToSubstraitType.toNamedStruct(agg.schema).names()
 
     relation.Project.builder
       .remap(relation.Rel.Remap.offset(newOutput.size, projects.size))
       .expressions(projects.asJava)
+      .hint(Hint.builder.addAllOutputNames(names).build())
       .input(substraitAgg)
       .build()
   }
@@ -340,9 +342,11 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging {
       p.child.output.count(o => !existenceJoins.contains(o.exprId.id)),
       expressions.size
     )
+
     relation.Project.builder
       .remap(remap)
       .expressions(expressions.asJava)
+      .hint(Hint.builder.addAllOutputNames(ToSubstraitType.toNamedStruct(p.schema).names()).build())
       .input(child)
       .build()
   }
@@ -357,11 +361,9 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging {
           .build()
       })
 
-    val names = p.output.map(_.name)
-
     relation.Expand.builder
       .fields(fields.asJava)
-      .hint(Hint.builder.addAllOutputNames(names.asJava).build())
+      .hint(Hint.builder.addAllOutputNames(ToSubstraitType.toNamedStruct(p.schema).names()).build())
       .input(visit(p.child))
       .build()
   }

spark/src/test/scala/io/substrait/spark/TPCDSPlan.scala

@@ -33,7 +33,6 @@ class TPCDSPlan extends TPCDSBase with SubstraitPlanTestBase {
 
   // spotless:off
   val failingSQL: Set[String] = Set(
-    "q35", "q51", "q84", // These fail when comparing the round-tripped query plans, but are actually equivalent (due to aliases being ignored by substrait)
     "q72" //requires implementation of date_add()
   )
   // spotless:on