[SPARK-16323] [SQL] Add IntegerDivide to avoid unnecessary cast

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -235,6 +235,7 @@ object FunctionRegistry {
     expression[Subtract]("-"),
     expression[Multiply]("*"),
     expression[Divide]("/"),
+    expression[IntegerDivide]("div"),
     expression[Remainder]("%"),
 
     // aggregate functions

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala

@@ -71,6 +71,7 @@ package object dsl {
     def - (other: Expression): Expression = Subtract(expr, other)
     def * (other: Expression): Expression = Multiply(expr, other)
     def / (other: Expression): Expression = Divide(expr, other)
+    def div (other: Expression): Expression = IntegerDivide(expr, other)
     def % (other: Expression): Expression = Remainder(expr, other)
     def & (other: Expression): Expression = BitwiseAnd(expr, other)
     def | (other: Expression): Expression = BitwiseOr(expr, other)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/arithmetic.scala

@@ -207,20 +207,12 @@ case class Multiply(left: Expression, right: Expression)
   protected override def nullSafeEval(input1: Any, input2: Any): Any = numeric.times(input1, input2)
 }
 
-@ExpressionDescription(
-  usage = "a _FUNC_ b - Divides a by b.",
-  extended = "> SELECT 3 _FUNC_ 2;\n 1.5")
-case class Divide(left: Expression, right: Expression)
-    extends BinaryArithmetic with NullIntolerant {
-
-  override def inputType: AbstractDataType = TypeCollection(DoubleType, DecimalType)
-
-  override def symbol: String = "/"
-  override def decimalMethod: String = "$div"
+abstract class DivisionArithmetic extends BinaryArithmetic with NullIntolerant {
   override def nullable: Boolean = true
 
   private lazy val div: (Any, Any) => Any = dataType match {
     case ft: FractionalType => ft.fractional.asInstanceOf[Fractional[Any]].div
+    case i: IntegralType => i.integral.asInstanceOf[Integral[Any]].quot
   }
 
   override def eval(input: InternalRow): Any = {
@@ -237,6 +229,9 @@ case class Divide(left: Expression, right: Expression)
     }
   }
 
+  // Used by doGenCode
+  protected def divide(eval1: ExprCode, eval2: ExprCode, javaType: String): String
+
   /**
    * Special case handling due to division by 0 => null.
    */
@@ -249,11 +244,7 @@ case class Divide(left: Expression, right: Expression)
       s"${eval2.value} == 0"
     }
     val javaType = ctx.javaType(dataType)
-    val divide = if (dataType.isInstanceOf[DecimalType]) {
-      s"${eval1.value}.$decimalMethod(${eval2.value})"
-    } else {
-      s"($javaType)(${eval1.value} $symbol ${eval2.value})"
-    }
+    val division = divide(eval1, eval2, javaType)
     if (!left.nullable && !right.nullable) {
       ev.copy(code = s"""
         ${eval2.code}
@@ -263,7 +254,7 @@ case class Divide(left: Expression, right: Expression)
           ${ev.isNull} = true;
         } else {
           ${eval1.code}
-          ${ev.value} = $divide;
+          ${ev.value} = $division;
         }""")
     } else {
       ev.copy(code = s"""
@@ -277,13 +268,51 @@ case class Divide(left: Expression, right: Expression)
           if (${eval1.isNull}) {
             ${ev.isNull} = true;
           } else {
-            ${ev.value} = $divide;
+            ${ev.value} = $division;
           }
         }""")
     }
   }
 }
 
+@ExpressionDescription(
+  usage = "a _FUNC_ b - Divides a by b.",
+  extended = "> SELECT 3 _FUNC_ 2;\n 1.5")
+case class Divide(left: Expression, right: Expression)
+    extends DivisionArithmetic {
+
+  override def inputType: AbstractDataType = TypeCollection(DoubleType, DecimalType)
+
+  override def symbol: String = "/"
+  override def decimalMethod: String = "$div"
+
+  // Used by doGenCode
+  protected override def divide(eval1: ExprCode, eval2: ExprCode, javaType: String): String = {
+    if (dataType.isInstanceOf[DecimalType]) {
+      s"${eval1.value}.$decimalMethod(${eval2.value})"
+    } else {
+      s"($javaType)(${eval1.value} $symbol ${eval2.value})"
+    }
+  }
+}
+
+@ExpressionDescription(
+  usage = "a _FUNC_ b - Divides a by b.",
+  extended = "> SELECT 3 _FUNC_ 2;\n 1")
+case class IntegerDivide(left: Expression, right: Expression)
+  extends DivisionArithmetic {
+
+  override def inputType: AbstractDataType = IntegralType
+
+  override def symbol: String = "div"
+  override def decimalMethod: String = "/"
+
+  // Used by doGenCode
+  protected override def divide(eval1: ExprCode, eval2: ExprCode, javaType: String): String = {
+    s"($javaType)(${eval1.value} $decimalMethod (${eval2.value}))"
+  }
+}
+
 @ExpressionDescription(
   usage = "a _FUNC_ b - Returns the remainder when dividing a by b.")
 case class Remainder(left: Expression, right: Expression)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -957,7 +957,7 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with Logging {
       case SqlBaseParser.PERCENT =>
         Remainder(left, right)
       case SqlBaseParser.DIV =>
-        Cast(Divide(left, right), LongType)
+        IntegerDivide(left, right)
       case SqlBaseParser.PLUS =>
         Add(left, right)
       case SqlBaseParser.MINUS =>

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ArithmeticExpressionSuite.scala

@@ -140,14 +140,14 @@ class ArithmeticExpressionSuite extends SparkFunSuite with ExpressionEvalHelper
 
   // By fixing SPARK-15776, Divide's inputType is required to be DoubleType of DecimalType.
   // TODO: in future release, we should add a IntegerDivide to support integral types.
-  ignore("/ (Divide) for integral type") {
-    checkEvaluation(Divide(Literal(1.toByte), Literal(2.toByte)), 0.toByte)
-    checkEvaluation(Divide(Literal(1.toShort), Literal(2.toShort)), 0.toShort)
-    checkEvaluation(Divide(Literal(1), Literal(2)), 0)
-    checkEvaluation(Divide(Literal(1.toLong), Literal(2.toLong)), 0.toLong)
-    checkEvaluation(Divide(positiveShortLit, negativeShortLit), 0.toShort)
-    checkEvaluation(Divide(positiveIntLit, negativeIntLit), 0)
-    checkEvaluation(Divide(positiveLongLit, negativeLongLit), 0L)
+  test("/ (Divide) for integral type") {
+    checkEvaluation(IntegerDivide(Literal(1.toByte), Literal(2.toByte)), 0.toByte)
+    checkEvaluation(IntegerDivide(Literal(1.toShort), Literal(2.toShort)), 0.toShort)
+    checkEvaluation(IntegerDivide(Literal(1), Literal(2)), 0)
+    checkEvaluation(IntegerDivide(Literal(1.toLong), Literal(2.toLong)), 0.toLong)
+    checkEvaluation(IntegerDivide(positiveShortLit, negativeShortLit), 0.toShort)
+    checkEvaluation(IntegerDivide(positiveIntLit, negativeIntLit), 0)
+    checkEvaluation(IntegerDivide(positiveLongLit, negativeLongLit), 0L)
   }
 
   test("% (Remainder)") {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/ExpressionParserSuite.scala

@@ -169,7 +169,7 @@ class ExpressionParserSuite extends PlanTest {
     // Simple operations
     assertEqual("a * b", 'a * 'b)
     assertEqual("a / b", 'a / 'b)
-    assertEqual("a DIV b", ('a / 'b).cast(LongType))
+    assertEqual("a DIV b", ('a div 'b))
     assertEqual("a % b", 'a % 'b)
     assertEqual("a + b", 'a + 'b)
     assertEqual("a - b", 'a - 'b)
@@ -180,7 +180,7 @@ class ExpressionParserSuite extends PlanTest {
     // Check precedences
     assertEqual(
       "a * t | b ^ c & d - e + f % g DIV h / i * k",
-      'a * 't | ('b ^ ('c & ('d - 'e + (('f % 'g / 'h).cast(LongType) / 'i * 'k)))))
+      'a * 't | ('b ^ ('c & ('d - 'e + (('f % 'g div 'h) / 'i * 'k)))))
   }
 
   test("unary arithmetic expressions") {

sql/core/src/main/scala/org/apache/spark/sql/Column.scala

@@ -731,6 +731,21 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    */
   def / (other: Any): Column = withExpr { Divide(expr, lit(other).expr) }
 
+  /**
+   * Integer Division this expression by another expression.
+   * {{{
+   *   // Scala: The following divides a person's height by their weight.
+   *   people.select( people("height") div people("weight") )
+   *
+   *   // Java:
+   *   people.select( people("height").div(people("weight")) );
+   * }}}
+   *
+   * @group expr_ops
+   * @since 2.1.0
+   */
+  def div (other: Any): Column = withExpr { IntegerDivide(expr, lit(other).expr) }
+
   /**
    * Division this expression by another expression.
    * {{{