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

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

@@ -234,6 +234,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

@@ -216,7 +216,6 @@ case class Divide(left: Expression, right: Expression)
   override def inputType: AbstractDataType = TypeCollection(DoubleType, DecimalType)
 
   override def symbol: String = "/"
-  override def decimalMethod: String = "$div"
   override def nullable: Boolean = true
 
   private lazy val div: (Any, Any) => Any = dataType match {
@@ -284,6 +283,75 @@ case class Divide(left: Expression, right: Expression)
   }
 }
 
+@ExpressionDescription(
+  usage = "a _FUNC_ b - Divides a by b.",
+  extended = "> SELECT 3 _FUNC_ 2;\n 1")
+case class IntegerDivide(left: Expression, right: Expression)
+  extends BinaryArithmetic with NullIntolerant {
+
+  override def inputType: AbstractDataType = IntegralType
+
+  override def symbol: String = "/"
+  override def decimalMethod: String = "$div"
+  override def nullable: Boolean = true
+
+  private lazy val div: (Any, Any) => Any = dataType match {
+    case i: IntegralType => i.integral.asInstanceOf[Integral[Any]].quot
+  }
+
+  override def eval(input: InternalRow): Any = {
+    val input2 = right.eval(input)
+    if (input2 == null || input2 == 0) {
+      null
+    } else {
+      val input1 = left.eval(input)
+      if (input1 == null) {
+        null
+      } else {
+        div(input1, input2)
+      }
+    }
+  }
+
+  /**
+   * Special case handling due to division by 0 => null.
+   */
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val eval1 = left.genCode(ctx)
+    val eval2 = right.genCode(ctx)
+    val isZero = s"${eval2.value} == 0"
+    val javaType = ctx.javaType(dataType)
+    val divide = s"($javaType)(${eval1.value} $symbol ${eval2.value})"
+    if (!left.nullable && !right.nullable) {
+      ev.copy(code = s"""
+        ${eval2.code}
+        boolean ${ev.isNull} = false;
+        $javaType ${ev.value} = ${ctx.defaultValue(javaType)};
+        if ($isZero) {
+          ${ev.isNull} = true;
+        } else {
+          ${eval1.code}
+          ${ev.value} = $divide;
+        }""")
+    } else {
+      ev.copy(code = s"""
+        ${eval2.code}
+        boolean ${ev.isNull} = false;
+        $javaType ${ev.value} = ${ctx.defaultValue(javaType)};
+        if (${eval2.isNull} || $isZero) {
+          ${ev.isNull} = true;
+        } else {
+          ${eval1.code}
+          if (${eval1.isNull}) {
+            ${ev.isNull} = true;
+          } else {
+            ${ev.value} = $divide;
+          }
+        }""")
+    }
+  }
+}
+
 @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.
    * {{{