Merge pull request #1279 from datastax/SPARKC-613-2.5

SPARKC-613 fix direct join for spark sql

Author: jtgrabowski

connector/src/it/scala/org/apache/spark/sql/cassandra/execution/CassandraDirectJoinSpec.scala

@@ -2,7 +2,6 @@ package org.apache.spark.sql.cassandra.execution
 
 import java.sql.Timestamp
 import java.time.Instant
-import java.util.concurrent.CompletableFuture
 
 import com.datastax.oss.driver.api.core.DefaultProtocolVersion._
 import com.datastax.spark.connector.cluster.DefaultCluster
@@ -634,6 +633,18 @@ class CassandraDirectJoinSpec extends SparkCassandraITFlatSpecBase with DefaultC
     quotes.count should be (1)
   }
 
+  /** SPARKC-613 */
+  it should "use direct join for sql join" in compareDirectOnDirectOff { spark =>
+    import spark.implicits._
+
+    val toJoin = spark.range(1, 5).map(_.intValue).withColumnRenamed("value", "id")
+
+    toJoin.createOrReplaceTempView("tojoin")
+    spark.read.cassandraFormat("kv", ks).load().createOrReplaceTempView(s"cassdata")
+
+    spark.sql("select * from tojoin tj inner join cassdata cd on tj.id = cd.k")
+  }
+
   private def compareDirectOnDirectOff(test: ((SparkSession) => DataFrame)) = {
     val sparkJoinOn = sparkSession.cloneSession()
     sparkJoinOn.conf.set(DirectJoinSettingParam.name, "on")

connector/src/main/scala/org/apache/spark/sql/cassandra/execution/CassandraDirectJoinExec.scala

@@ -5,7 +5,7 @@ import com.datastax.spark.connector.rdd.{CassandraJoinRDD, CassandraLeftJoinRDD,
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.cassandra.execution.unsafe.{UnsafeRowReaderFactory, UnsafeRowWriterFactory}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BindReferences, EqualTo, Expression, GenericInternalRow, JoinedRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BindReferences, EqualTo, ExprId, Expression, GenericInternalRow, JoinedRow, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.execution.{DataSourceScanExec, SparkPlan, UnaryExecNode}
 import org.apache.spark.sql.execution.joins.{BuildLeft, BuildSide}
@@ -21,7 +21,7 @@ case class CassandraDirectJoinExec(
   cassandraSide: BuildSide,
   condition: Option[Expression],
   child: SparkPlan,
-  aliasMap: Map[String, Attribute],
+  aliasMap: Map[String, ExprId],
   cassandraScan: CassandraTableScanRDD[_],
   cassandraPlan: DataSourceScanExec) extends UnaryExecNode {
 
@@ -41,7 +41,7 @@ case class CassandraDirectJoinExec(
   val primaryKeys = cassandraScan.tableDef.primaryKey.map(_.columnName)
   val cassandraSchema = cassandraPlan.schema
 
-  val attributeToCassandra = aliasMap.map(_.swap)
+  val exprIdToCassandra = aliasMap.map(_.swap)
 
   val leftJoinCouplets =
     if (cassandraSide == BuildLeft) leftKeys.zip(rightKeys) else rightKeys.zip(leftKeys)
@@ -54,15 +54,15 @@ case class CassandraDirectJoinExec(
     */
   val (pkJoinCoulplets, otherJoinCouplets) = leftJoinCouplets.partition {
     case (cassandraAttribute: Attribute, _) =>
-      attributeToCassandra.get(cassandraAttribute) match {
+      exprIdToCassandra.get(cassandraAttribute.exprId) match {
         case Some(name) if primaryKeys.contains(name) => true
         case _ => false
-    }
+      }
     case _ => false
   }
 
   val (joinColumns, joinExpressions) = pkJoinCoulplets.map { case (cAttr: Attribute, otherCol: Expression) =>
-    (ColumnName(attributeToCassandra(cAttr)), BindReferences.bindReference(otherCol, keySource.output))
+    (ColumnName(exprIdToCassandra(cAttr.exprId)), BindReferences.bindReference(otherCol, keySource.output))
   }.unzip
 
   /**
@@ -210,7 +210,7 @@ case class CassandraDirectJoinExec(
     val selectString = selectedColumns.mkString("Reading (", ", ", ")")
 
     val joinString = pkJoinCoulplets
-      .map{ case (colref: Attribute, exp) => s"${attributeToCassandra(colref)} = ${exp}"}
+      .map{ case (colref: Attribute, exp) => s"${exprIdToCassandra(colref.exprId)} = ${exp}"}
       .mkString(", ")
 
     s"Cassandra Direct Join [${joinString}] $keyspace.$table - $selectString${pushedWhere} "

connector/src/main/scala/org/apache/spark/sql/cassandra/execution/CassandraDirectJoinStrategy.scala

@@ -298,12 +298,14 @@ object CassandraDirectJoinStrategy extends Logging {
   def allPartitionKeysAreJoined(plan: LogicalPlan, joinKeys: Seq[Expression]): Boolean =
     plan match {
       case PhysicalOperation(
-        attributes, _,
-        LogicalRelation(cassandraSource: CassandraSourceRelation, _, _, _)) =>
+      attributes, _,
+      LogicalRelation(cassandraSource: CassandraSourceRelation, _, _, _)) =>
+
+        val joinKeysExprId = joinKeys.collect { case attributeReference: AttributeReference => attributeReference.exprId }
 
         val joinKeyAliases =
           aliasMap(attributes)
-            .filter{ case (_, value) => joinKeys.contains(value) }
+            .filter { case (_, value) => joinKeysExprId.contains(value) }
         val partitionKeyNames = cassandraSource.tableDef.partitionKey.map(_.columnName)
         val allKeysPresent = partitionKeyNames.forall(joinKeyAliases.contains)
 
@@ -312,15 +314,15 @@ object CassandraDirectJoinStrategy extends Logging {
         }
 
         allKeysPresent
-    case _ => false
-  }
+      case _ => false
+    }
 
   /**
     * Map Source Names to Attributes
     */
   def aliasMap(aliases: Seq[NamedExpression]) = aliases.map {
-    case a @ Alias(child: AttributeReference, _) => child.name -> a.toAttribute
-    case namedExpression: NamedExpression => namedExpression.name -> namedExpression.toAttribute
+    case a @ Alias(child: AttributeReference, _) => child.name -> a.exprId
+    case attributeReference: AttributeReference => attributeReference.name -> attributeReference.exprId
   }.toMap
 
   /**
@@ -335,6 +337,4 @@ object CassandraDirectJoinStrategy extends Logging {
     }
   }
 
-
-
 }