org.apache.spark.sql.cassandra.PredicateOps.scala Maven / Gradle / Ivy

package org.apache.spark.sql.cassandra

import org.apache.spark.sql.catalyst.expressions
import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.sources
import org.apache.spark.sql.sources.Filter

/** A unified API for predicates, used by [[BasicCassandraPredicatePushDown]].
  *
  * Keeps all the Spark-specific stuff out of `BasicCassandraPredicatePushDown`
  * It is also easy to plug-in custom predicate implementations for unit-testing.
  */
trait PredicateOps[Predicate] {

  /** Returns the column name of the only column in the predicate.
    * Will throw `IllegalArgumentException` if predicate is not a single-column predicate. */
  def columnName(p: Predicate): String

  /** Returns true if the predicate restricts exactly one column. */
  def isSingleColumnPredicate(p: Predicate): Boolean

  /** Returns true for predicates of type: column = value */
  def isEqualToPredicate(p: Predicate): Boolean

  /** Returns true for predicates of type: column (< | <= | >= | >) value */
  def isRangePredicate(p: Predicate): Boolean

  /** Returns true for predicates of type: column IN (value1, value2, ...) */
  def isInPredicate(p: Predicate): Boolean
}

/** Provides `PredicateOps` adapters for Expression and Filter classes */
object PredicateOps {

  /** Adapts the API of Catalyst `Expression` to the API expected by `BasicCassandraPredicatePushDown` */
  implicit object ExpressionOps extends PredicateOps[Expression] {

    override def columnName(p: Expression): String = {
      require(isSingleColumnPredicate(p), s"Not a single-column predicate: $p")
      p.references.head.name
    }

    override def isSingleColumnPredicate(p: Expression): Boolean =
      p.references.size == 1

    override def isRangePredicate(p: Expression): Boolean = p match {
      case _: expressions.LessThan => true
      case _: expressions.LessThanOrEqual => true
      case _: expressions.GreaterThan => true
      case _: expressions.GreaterThanOrEqual => true
      case _ => false
    }

    override def isEqualToPredicate(p: Expression): Boolean =
      p.isInstanceOf[expressions.EqualTo]

    override def isInPredicate(p: Expression): Boolean =
      p.isInstanceOf[expressions.In] || p.isInstanceOf[expressions.InSet]

  }

  /** Adapts the API of Catalyst `Filter` to the API expected by `BasicCassandraPredicatePushDown` */
  implicit object FilterOps extends PredicateOps[Filter] {

    override def columnName(p: Filter): String = p match {
      case eq: sources.EqualTo => eq.attribute
      case lt: sources.LessThan => lt.attribute
      case lte: sources.LessThanOrEqual => lte.attribute
      case gt: sources.GreaterThan => gt.attribute
      case gte: sources.GreaterThanOrEqual => gte.attribute
      case in: sources.In => in.attribute
      case _ => throw new IllegalArgumentException(
        s"Don't know how to get column name from the predicate: $p")
    }

    override def isSingleColumnPredicate(p: Filter): Boolean =
      isRangePredicate(p) || isEqualToPredicate(p) || isInPredicate(p)

    override def isRangePredicate(p: Filter): Boolean = p match {
      case _: sources.LessThan => true
      case _: sources.LessThanOrEqual => true
      case _: sources.GreaterThan => true
      case _: sources.GreaterThanOrEqual => true
      case _ => false
    }

    override def isEqualToPredicate(p: Filter): Boolean =
      p.isInstanceOf[sources.EqualTo]

    override def isInPredicate(p: Filter): Boolean =
      p.isInstanceOf[sources.In]
  }
}