[SPARK-15453] [SQL] FileSourceScanExec to extract outputOrdering information

sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala

@@ -23,7 +23,7 @@ import org.apache.commons.lang3.StringUtils
 import org.apache.hadoop.fs.{BlockLocation, FileStatus, LocatedFileStatus, Path}
 
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{Row, SparkSession, SQLContext}
+import org.apache.spark.sql.{AnalysisException, Row, SparkSession, SQLContext}
 import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier}
 import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.expressions._
@@ -156,24 +156,56 @@ case class FileSourceScanExec(
     false
   }
 
-  override val outputPartitioning: Partitioning = {
+  override val (outputPartitioning, outputOrdering): (Partitioning, Seq[SortOrder]) = {
     val bucketSpec = if (relation.sparkSession.sessionState.conf.bucketingEnabled) {
       relation.bucketSpec
     } else {
       None
     }
-    bucketSpec.map { spec =>
-      val numBuckets = spec.numBuckets
-      val bucketColumns = spec.bucketColumnNames.flatMap { n =>
-        output.find(_.name == n)
-      }
-      if (bucketColumns.size == spec.bucketColumnNames.size) {
-        HashPartitioning(bucketColumns, numBuckets)
-      } else {
-        UnknownPartitioning(0)
-      }
-    }.getOrElse {
-      UnknownPartitioning(0)
+    bucketSpec match {
+      case Some(spec) =>
+        val numBuckets = spec.numBuckets
+        val bucketColumns = spec.bucketColumnNames.flatMap { n =>
+          output.find(_.name == n)
+        }
+        if (bucketColumns.size == spec.bucketColumnNames.size) {
+          val partitioning = HashPartitioning(bucketColumns, numBuckets)
+
+          val sortOrder = if (spec.sortColumnNames.nonEmpty) {
+            // In case of bucketing, its possible to have multiple files belonging to the
+            // same bucket in a given relation. Each of these files are locally sorted
+            // but those files combined together are not globally sorted. Given that,
+            // the RDD partition will not be sorted even if the relation has sort columns set
+            // Current solution is to check if all the buckets have a single file in it
+
+            val files =
+              relation.location.listFiles(partitionFilters).flatMap(partition => partition.files)
+            val bucketToFilesGrouping =
+              files.map(_.getPath.getName).groupBy(file => BucketingUtils.getBucketId(file))
+            val singleFilePartitions = bucketToFilesGrouping.forall(p => p._2.length <= 1)
+
+            if (singleFilePartitions) {
+              def toAttribute(colName: String): Attribute =
+                output.find(_.name == colName).getOrElse {
+                  throw new AnalysisException(s"Could not find sort column $colName for " +
+                    s"relation ${metastoreTableIdentifier.get.toString} in its existing " +
+                    s"columns : (${output.map(_.name).mkString(", ")})")
+              }
+              // TODO Currently Spark does not support writing columns sorting in descending order
+              // so using Ascending order. This can be fixed in future
+              spec.sortColumnNames.map(c => SortOrder(toAttribute(c), Ascending))
+            } else {
+              Nil
+            }
+          } else {
+            Nil
+          }
+          (partitioning, sortOrder)
+        } else {
+          (UnknownPartitioning(0), Nil)
+        }
+      case _ =>
+        (UnknownPartitioning(0), Nil)
     }
   }
 

sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala

@@ -21,6 +21,7 @@ import scala.language.existentials
 
 import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
 import org.apache.spark.sql.catalyst.TableIdentifier
+import org.apache.spark.sql.execution.SortExec
 import org.apache.spark.sql.execution.joins._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSQLContext
@@ -61,6 +62,51 @@ class JoinSuite extends QueryTest with SharedSQLContext {
     }
   }
 
+  test("SPARK-15453 : Sort Merge join on bucketed + sorted tables should not add `sort` step " +
+    "if the join predicates are subset of the sorted columns of the tables") {
+    withTable("SPARK_15453_table_a", "SPARK_15453_table_b") {
+      withSQLConf("spark.sql.autoBroadcastJoinThreshold" -> "0") {
+        val df =
+          (0 until 8)
+            .map(i => (i, i * 2, i.toString))
+            .toDF("i", "j", "k")
+            .coalesce(1)
+        df.write.bucketBy(4, "j", "k").sortBy("j", "k").saveAsTable("SPARK_15453_table_a")
+        df.write.bucketBy(4, "j", "k").sortBy("j", "k").saveAsTable("SPARK_15453_table_b")
+
+        val query = """
+                      |SELECT *
+                      |FROM
+                      |  SPARK_15453_table_a a
+                      |JOIN
+                      |  SPARK_15453_table_b b
+                      |ON a.j=b.j AND
+                      |   a.k=b.k
+                    """.stripMargin
+        val joinDF = sql(query)
+
+        val executedPlan = joinDF.queryExecution.executedPlan
+        val operators = executedPlan.collect {
+          case j: SortMergeJoinExec => j
+          case j: SortExec => j
+        }
+        assert(operators.size === 1)
+        assert(operators.head.getClass == classOf[SortMergeJoinExec])
+
+        checkAnswer(joinDF,
+          Row(0, 0, "0", 0, 0, "0") ::
+            Row(1, 2, "1", 1, 2, "1") ::
+            Row(2, 4, "2", 2, 4, "2") ::
+            Row(3, 6, "3", 3, 6, "3") ::
+            Row(4, 8, "4", 4, 8, "4") ::
+            Row(5, 10, "5", 5, 10, "5") ::
+            Row(6, 12, "6", 6, 12, "6") ::
+            Row(7, 14, "7", 7, 14, "7") :: Nil)
+      }
+    }
+  }
+
+
   test("join operator selection") {
     spark.sharedState.cacheManager.clearCache()