[SPARK-19931][SQL] InMemoryTableScanExec should rewrite output partitioning and ordering when aliasing output attributes

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala

@@ -22,7 +22,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.QueryPlan
-import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning}
 import org.apache.spark.sql.execution.LeafExecNode
 import org.apache.spark.sql.execution.metric.SQLMetrics
 import org.apache.spark.sql.types.UserDefinedType
@@ -42,10 +42,34 @@ case class InMemoryTableScanExec(
   override def output: Seq[Attribute] = attributes
 
   // The cached version does not change the outputPartitioning of the original SparkPlan.
-  override def outputPartitioning: Partitioning = relation.child.outputPartitioning
+  // But the cached version could alias output, so we need to replace output.
+  override def outputPartitioning: Partitioning = {
+    val attrMap = AttributeMap(
+      relation.child.output.zip(output)
+    )
+    relation.child.outputPartitioning match {
+      case HashPartitioning(expressions, numPartitions) =>
+        val newExprs = expressions.map(_.transform {
+          case attr: Attribute if attrMap.contains(attr) => attrMap.get(attr).get
+        })
+        HashPartitioning(newExprs, numPartitions)
+      case _ => relation.child.outputPartitioning
+    }
+  }
 
   // The cached version does not change the outputOrdering of the original SparkPlan.
-  override def outputOrdering: Seq[SortOrder] = relation.child.outputOrdering
+  // But the cached version could alias output, so we need to replace output.
+  override def outputOrdering: Seq[SortOrder] = {
+    val attrMap = AttributeMap(
+      relation.child.output.zip(output)
+    )
+    relation.child.outputOrdering.map { sortOrder =>
+      val newSortExpr = sortOrder.child.transform {
+        case attr: Attribute if attrMap.contains(attr) => attrMap.get(attr).get
+      }
+      SortOrder(newSortExpr, sortOrder.direction, sortOrder.nullOrdering)
+    }
+  }
 
   private def statsFor(a: Attribute) = relation.partitionStatistics.forAttribute(a)
 

sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala

@@ -21,6 +21,9 @@ import java.nio.charset.StandardCharsets
 import java.sql.{Date, Timestamp}
 
 import org.apache.spark.sql.{DataFrame, QueryTest, Row}
+import org.apache.spark.sql.catalyst.expressions.AttributeSet
+import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
+import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.test.SQLTestData._
@@ -390,4 +393,23 @@ class InMemoryColumnarQuerySuite extends QueryTest with SharedSQLContext {
     }
   }
 
+  test("InMemoryTableScanExec should return currect output ordering and partitioning") {
+    val ds1 = Seq((0, 0), (1, 1)).toDS
+      .repartition(col("_1")).sortWithinPartitions(col("_1")).persist
+    val ds2 = Seq((0, 0), (1, 1)).toDS
+      .repartition(col("_1")).sortWithinPartitions(col("_1")).persist
+    val joined = ds1.joinWith(ds2, ds1("_1") === ds2("_1"))
+
+    val inMemoryScans = joined.queryExecution.executedPlan.collect {
+      case m: InMemoryTableScanExec => m
+    }
+    inMemoryScans.foreach { inMemoryScan =>
+      val sortedAttrs = AttributeSet(inMemoryScan.outputOrdering.flatMap(_.references))
+      assert(sortedAttrs.subsetOf(inMemoryScan.outputSet))
+
+      val partitionedAttrs =
+        inMemoryScan.outputPartitioning.asInstanceOf[HashPartitioning].references
+      assert(partitionedAttrs.subsetOf(inMemoryScan.outputSet))
+    }
+  }
 }