SPARKC-621 move timeuuid non-equal predicates check at the end

Before this commit the check was executed in
BasicCassandraPredicatePushDown at the beginning of the push down
evaluation. This interrupted push down evaluation without giving
other predicates rules a chance to push down the predicate.
Since thanks to SAI we may now push timeuuid predicates, this
commit moves the check at the very end of the push down evaluation.

Author: jtgrabowski

connector/src/it/scala/com/datastax/spark/connector/SparkCassandraITFlatSpecBase.scala

@@ -187,6 +187,12 @@ trait SparkCassandraITSpecBase
     else report(s"Skipped because not DSE")
   }
 
+  /** Skips the given test if the cluster is not Cassandra */
+  def cassandraOnly(f: => Unit): Unit = {
+    if (isDse(conn)) report(s"Skipped because not Cassandra")
+    else f
+  }
+
   /** Skips the given test if the Cluster Version is lower or equal to the given version or the cluster is not DSE */
   def dseFrom(version: Version)(f: => Any): Unit = {
     dseOnly {

connector/src/it/scala/com/datastax/spark/connector/sql/CassandraDataFrameSpec.scala

@@ -74,7 +74,14 @@ class CassandraDataFrameSpec extends SparkCassandraITFlatSpecBase with DefaultCl
             session.execute(s"create table $ks.date_test2 (key int primary key, dd date)")
             session.execute(s"insert into $ks.date_test (key, dd) values (1, '1930-05-31')")
           }
-        }
+        },
+
+        Future {
+          session.execute(
+            s"""
+               |CREATE TABLE $ks.timeuuidtable (k INT, v TIMEUUID, PRIMARY KEY (k))
+               |""".stripMargin)
+        },
       )
       executor.waitForCurrentlyExecutingTasks()
     }
@@ -356,5 +363,23 @@ class CassandraDataFrameSpec extends SparkCassandraITFlatSpecBase with DefaultCl
 
   }
 
+  it should "complain when non-equality predicate on timeuuid is detected" in cassandraOnly {
+    val df = sparkSession
+      .read
+      .format("org.apache.spark.sql.cassandra")
+      .options(
+        Map(
+          "table" -> "timeuuidtable",
+          "keyspace" -> ks
+        )
+      )
+      .load()
+      .where(col("v") < "61129592-FBE4-11E3-A3AC-0800200C9A66")
+
+
+    intercept[IllegalArgumentException] {
+      df.count()
+    }
+  }
 
 }

connector/src/main/scala/org/apache/spark/sql/cassandra/BasicCassandraPredicatePushDown.scala

@@ -69,20 +69,6 @@ class BasicCassandraPredicatePushDown[Predicate : PredicateOps](
   private def firstNonEmptySet[T](sets: Set[T]*): Set[T] =
     sets.find(_.nonEmpty).getOrElse(Set.empty[T])
 
-
-  /** All non-equal predicates on a TimeUUID column are going to fail and fail
-    * in silent way. The basic issue here is that when you use a comparison on
-    * a time UUID column in C* it compares based on the Time portion of the UUID. When
-    * Spark executes this filter (unhandled behavior) it will compare lexically, this
-    * will lead to results being incorrectly filtered out of the set. As long as the
-    * range predicate is handled completely by the connector the correct result
-    * will be obtained.
-    */
-  val timeUUIDNonEqual = {
-    val timeUUIDCols = table.columns.filter(x => x.columnType == TimeUUIDType)
-    timeUUIDCols.flatMap(col => rangePredicatesByName.get(col.columnName)).flatten
-  }
-
   /** Evaluates set of columns used in equality predicates and set of columns use in 'in' predicates.
     * Evaluation is based on protocol version. */
   private def partitionKeyPredicateColumns(pv: ProtocolVersion): (Seq[String], Seq[String]) = {
@@ -197,20 +183,4 @@ class BasicCassandraPredicatePushDown[Predicate : PredicateOps](
   val predicatesToPreserve: Set[Predicate] =
     predicates -- predicatesToPushDown
 
-
-  val unhandledTimeUUIDNonEqual = {
-    timeUUIDNonEqual.toSet -- predicatesToPushDown
-  }
-
-  require(unhandledTimeUUIDNonEqual.isEmpty,
-    s"""
-      | You are attempting to do a non-equality comparison on a TimeUUID column in Spark.
-      | Spark can only compare TimeUUIDs Lexically which means that the comparison will be
-      | different than the comparison done in C* which is done based on the Time Portion of
-      | TimeUUID. This will in almost all cases lead to incorrect results. If possible restrict
-      | doing a TimeUUID comparison only to columns which can be pushed down to Cassandra.
-      | https://datastax-oss.atlassian.net/browse/SPARKC-405.
-      |
-      | $unhandledTimeUUIDNonEqual
-    """.stripMargin)
 }

connector/src/main/scala/org/apache/spark/sql/cassandra/CassandraSourceRelation.scala

@@ -279,7 +279,8 @@ case class CassandraSourceRelation(
       DsePredicateRules,
       InClausePredicateRules) ++
       solrPredicateRules ++
-      additionalRules
+      additionalRules :+
+      TimeUUIDPredicateRules
 
     /** Apply non-basic rules **/
     val finalPushdown = predicatePushDownRules.foldLeft(basicPushdown)(

connector/src/main/scala/org/apache/spark/sql/cassandra/TimeUUIDPredicateRules.scala

@@ -0,0 +1,54 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Please see the included license file for details.
+ */
+
+package org.apache.spark.sql.cassandra
+
+import com.datastax.spark.connector.cql.TableDef
+import com.datastax.spark.connector.types.TimeUUIDType
+import com.datastax.spark.connector.util.Logging
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.cassandra.PredicateOps.FilterOps
+import org.apache.spark.sql.sources.Filter
+
+/** All non-equal predicates on a TimeUUID column are going to fail and fail
+  * in silent way. The basic issue here is that when you use a comparison on
+  * a time UUID column in C* it compares based on the Time portion of the UUID. When
+  * Spark executes this filter (unhandled behavior) it will compare lexically, this
+  * will lead to results being incorrectly filtered out of the set. As long as the
+  * range predicate is handled completely by the connector the correct result
+  * will be obtained.
+  */
+object TimeUUIDPredicateRules extends CassandraPredicateRules with Logging {
+
+  private def isTimeUUIDNonEqualPredicate(tableDef: TableDef, predicate: Filter): Boolean = {
+    if (FilterOps.isSingleColumnPredicate(predicate)) {
+      // extract column name only from single column predicates, otherwise an exception is thrown
+      val columnName = FilterOps.columnName(predicate)
+      val isRange = FilterOps.isRangePredicate(predicate)
+      val isTimeUUID = tableDef.columnByName.get(columnName).exists(_.columnType == TimeUUIDType)
+      isRange && isTimeUUID
+    } else {
+      false
+    }
+  }
+
+  override def apply(predicates: AnalyzedPredicates, tableDef: TableDef, conf: SparkConf): AnalyzedPredicates = {
+    val unhandledTimeUUIDNonEqual = predicates.handledBySpark.filter(isTimeUUIDNonEqualPredicate(tableDef, _))
+    require(unhandledTimeUUIDNonEqual.isEmpty,
+      s"""
+         | You are attempting to do a non-equality comparison on a TimeUUID column in Spark.
+         | Spark can only compare TimeUUIDs Lexically which means that the comparison will be
+         | different than the comparison done in C* which is done based on the Time Portion of
+         | TimeUUID. This will in almost all cases lead to incorrect results. If possible restrict
+         | doing a TimeUUID comparison only to columns which can be pushed down to Cassandra.
+         | https://datastax-oss.atlassian.net/browse/SPARKC-405.
+         |
+         | $unhandledTimeUUIDNonEqual
+    """.stripMargin)
+
+    predicates
+  }
+}

connector/src/test/scala/org/apache/spark/sql/cassandra/PredicatePushDownSpec.scala

@@ -82,14 +82,6 @@ class PredicatePushDownSpec extends FlatSpec with Matchers {
     ppd.predicatesToPreserve shouldBe empty
   }
 
-  it should " break if the user tries to use a TimeUUID on a fully unhandled predicate" in {
-    val f1 = GtFilter("t1")
-
-    val ex = intercept[IllegalArgumentException] {
-      val ppd = new BasicCassandraPredicatePushDown(Set[Filter](f1), timeUUIDTable)
-    }
-  }
-
   it should " work if the user tries to use a TimeUUID on a fully handled predicate" in {
     val f1 = GtFilter("c1")
     val ppd = new BasicCassandraPredicatePushDown(Set[Filter](f1), timeUUIDTable)