[SPARK-21117][SQL] Built-in SQL Function Support - WIDTH_BUCKET

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

@@ -262,6 +262,7 @@ object FunctionRegistry {
     expression[Tan]("tan"),
     expression[Cot]("cot"),
     expression[Tanh]("tanh"),
+    expression[WidthBucket]("width_bucket"),
 
     expression[Add]("+"),
     expression[Subtract]("-"),

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

@@ -48,6 +48,7 @@ import org.apache.spark.util.Utils
  * - [[UnaryExpression]]: an expression that has one child.
  * - [[BinaryExpression]]: an expression that has two children.
  * - [[TernaryExpression]]: an expression that has three children.
+ * - [[QuaternaryExpression]]: an expression that has four children.
  * - [[BinaryOperator]]: a special case of [[BinaryExpression]] that requires two children to have
  *                       the same output data type.
  *
@@ -631,3 +632,109 @@ abstract class TernaryExpression extends Expression {
     }
   }
 }
+
+/**
+ * An expression with four inputs and one output. The output is by default evaluated to null
+ * if any input is evaluated to null.
+ */
+abstract class QuaternaryExpression extends Expression {
+
+  override def foldable: Boolean = children.forall(_.foldable)
+
+  override def nullable: Boolean = children.exists(_.nullable)
+
+  /**
+   * Default behavior of evaluation according to the default nullability of QuaternaryExpression.
+   * If subclass of QuaternaryExpression override nullable, probably should also override this.
+   */
+  override def eval(input: InternalRow): Any = {
+    val exprs = children
+    val value1 = exprs(0).eval(input)
+    if (value1 != null) {
+      val value2 = exprs(1).eval(input)
+      if (value2 != null) {
+        val value3 = exprs(2).eval(input)
+        if (value3 != null) {
+          val value4 = exprs(3).eval(input)
+          if (value4 != null) {
+            return nullSafeEval(value1, value2, value3, value4)
+          }
+        }
+      }
+    }
+    null
+  }
+
+  /**
+   * Called by default [[eval]] implementation. If subclass of QuaternaryExpression keep the default
+   * nullability, they can override this method to save null-check code.  If we need full control
+   * of evaluation process, we should override [[eval]].
+   */
+  protected def nullSafeEval(input1: Any, input2: Any, input3: Any, input4: Any): Any =
+    sys.error(s"QuaternaryExpressions must override either eval or nullSafeEval")
+
+  /**
+   * Short hand for generating ternary evaluation code.
+   * If either of the sub-expressions is null, the result of this computation
+   * is assumed to be null.
+   *
+   * @param f accepts four variable names and returns Java code to compute the output.
+   */
+  protected def defineCodeGen(
+    ctx: CodegenContext,
+    ev: ExprCode,
+    f: (String, String, String, String) => String): ExprCode = {
+    nullSafeCodeGen(ctx, ev, (eval1, eval2, eval3, eval4) => {
+      s"${ev.value} = ${f(eval1, eval2, eval3, eval4)};"
+    })
+  }
+
+  /**
+   * Short hand for generating ternary evaluation code.
+   * If either of the sub-expressions is null, the result of this computation
+   * is assumed to be null.
+   *
+   * @param f function that accepts the 4 non-null evaluation result names of children
+   *          and returns Java code to compute the output.
+   */
+  protected def nullSafeCodeGen(
+    ctx: CodegenContext,
+    ev: ExprCode,
+    f: (String, String, String, String) => String): ExprCode = {
+    val firstGen = children(0).genCode(ctx)
+    val secondGen = children(1).genCode(ctx)
+    val thirdGen = children(2).genCode(ctx)
+    val fourthGen = children(3).genCode(ctx)
+    val resultCode = f(firstGen.value, secondGen.value, thirdGen.value, fourthGen.value)
+
+    if (nullable) {
+      val nullSafeEval =
+        firstGen.code + ctx.nullSafeExec(children(0).nullable, firstGen.isNull) {
+          secondGen.code + ctx.nullSafeExec(children(1).nullable, secondGen.isNull) {
+            thirdGen.code + ctx.nullSafeExec(children(2).nullable, thirdGen.isNull) {
+              fourthGen.code + ctx.nullSafeExec(children(3).nullable, fourthGen.isNull) {
+                s"""
+                  ${ev.isNull} = false; // resultCode could change nullability.
+                  $resultCode
+                """
+              }
+            }
+          }
+        }
+
+      ev.copy(code = s"""
+        boolean ${ev.isNull} = true;
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        $nullSafeEval""")
+    } else {
+      ev.copy(code = s"""
+        boolean ${ev.isNull} = false;
+        ${firstGen.code}
+        ${secondGen.code}
+        ${thirdGen.code}
+        ${fourthGen.code}
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        $resultCode""", isNull = "false")
+    }
+  }
+}

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

@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{TypeCheckFailure, TypeCheckSuccess}
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.util.NumberConverter
+import org.apache.spark.sql.catalyst.util.{MathUtils, NumberConverter}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
@@ -1186,3 +1186,51 @@ case class BRound(child: Expression, scale: Expression)
     with Serializable with ImplicitCastInputTypes {
   def this(child: Expression) = this(child, Literal(0))
 }
+
+/**
+ *  Returns the bucket number into which
+ *  the value of this expression would fall after being evaluated.
+ *
+ * @param expr is the expression for which the histogram is being created
+ * @param minValue is an expression that resolves
+ *                 to the minimum end point of the acceptable range for expr
+ * @param maxValue is an expression that resolves
+ *                 to the maximum end point of the acceptable range for expr
+ * @param numBucket is an An expression that resolves to
+ *                  a constant indicating the number of buckets
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+  usage = "_FUNC_(expr, min_value, max_value, num_bucket) - Returns an long between 0 and `num_buckets`+1 by mapping the `expr` into buckets defined by the range [`min_value`, `max_value`].",
+  extended = """
+    Examples:
+      > SELECT _FUNC_(5.35, 0.024, 10.06, 5);
+       3
+  """)
+// scalastyle:on line.size.limit
+case class WidthBucket(
+  expr: Expression,
+  minValue: Expression,
+  maxValue: Expression,
+  numBucket: Expression) extends QuaternaryExpression with ImplicitCastInputTypes {
+
+  override def children: Seq[Expression] = Seq(expr, minValue, maxValue, numBucket)
+  override def inputTypes: Seq[AbstractDataType] = Seq(DoubleType, DoubleType, DoubleType, LongType)
+  override def dataType: DataType = LongType
+  override def nullable: Boolean = true
+
+  override def nullSafeEval(ex: Any, min: Any, max: Any, num: Any): Any = {
+    MathUtils.widthBucket(
+      ex.asInstanceOf[Double],
+      min.asInstanceOf[Double],
+      max.asInstanceOf[Double],
+      num.asInstanceOf[Long])
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val mathUtils = MathUtils.getClass.getName.stripSuffix("$")
+    nullSafeCodeGen(ctx, ev, (ex, min, max, num) =>
+      s"${ev.value} = $mathUtils.widthBucket($ex, $min, $max, $num);"
+    )
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/MathUtils.scala

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.util
+
+import org.apache.spark.sql.AnalysisException
+
+object MathUtils {
+
+  /**
+   *  Returns the bucket number into which
+   *  the value of this expression would fall after being evaluated.
+   *
+   * @param expr id the expression for which the histogram is being created
+   * @param minValue is an expression that resolves
+   *                 to the minimum end point of the acceptable range for expr
+   * @param maxValue is an expression that resolves
+   *                 to the maximum end point of the acceptable range for expr
+   * @param numBucket is an An expression that resolves to
+   *                  a constant indicating the number of buckets
+   * @return Returns an long between 0 and numBucket+1 by mapping the expr into buckets defined by
+   *         the range [minValue, maxValue]. For example:
+   *         widthBucket(0, 1, 1, 1) -> 0, widthBucket(20, 1, 1, 1) -> 2.
+   */
+  def widthBucket(expr: Double, minValue: Double, maxValue: Double, numBucket: Long): Long = {
+
+    if (numBucket <= 0) {
+      throw new AnalysisException(s"The num of bucket must be greater than 0, but got ${numBucket}")
+    }
+
+    val lower: Double = Math.min(minValue, maxValue)
+    val upper: Double = Math.max(minValue, maxValue)
+
+    val result: Long = if (expr < lower) {
+      0
+    } else if (expr >= upper) {
+      numBucket + 1L
+    } else {
+      (numBucket.toDouble * (expr - lower) / (upper - lower) + 1).toLong
+    }
+
+    if (minValue > maxValue) (numBucket - result) + 1 else result
+  }
+}

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

@@ -22,6 +22,7 @@ import java.nio.charset.StandardCharsets
 import com.google.common.math.LongMath
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.TypeCoercion.ImplicitTypeCasts.implicitCast
 import org.apache.spark.sql.catalyst.dsl.expressions._
@@ -644,4 +645,36 @@ class MathExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(BRound(-0.35, 1), -0.4)
     checkEvaluation(BRound(-35, -1), -40)
   }
+
+  test("width_bucket") {
+    def test(
+      expr: Double,
+      minValue: Double,
+      maxValue: Double,
+      numBucket: Long,
+      expected: Long): Unit = {
+      checkEvaluation(WidthBucket(Literal.create(expr, DoubleType),
+        Literal.create(minValue, DoubleType),
+        Literal.create(maxValue, DoubleType),
+        Literal.create(numBucket, LongType)),
+        expected)
+    }
+
+    test(5.35, 0.024, 10.06, 5, 3)
+
+    test(3.14, 0, 4, 3, 3)
+    test(2, 0, 4, 3, 2)
+    test(-1, 0, 3.2, 4, 0)
+
+    test(3.14, 4, 0, 3, 1)
+    test(2, 4, 0, 3, 2)
+    test(-1, 3.2, 0, 4, 5)
+
+    intercept[AnalysisException]{
+      WidthBucket(Literal.create(1.0, DoubleType),
+        Literal.create(1.0, DoubleType),
+        Literal.create(2.0, DoubleType),
+        Literal.create(-1L, LongType)).eval()
+    }
+  }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/MathUtilsSuite.scala

@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.util
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.util.MathUtils._
+
+class MathUtilsSuite extends SparkFunSuite {
+
+  test("widthBucket") {
+    assert(widthBucket(5.35, 0.024, 10.06, 5) === 3)
+    assert(widthBucket(0, 1, 1, 1) === 0)
+    assert(widthBucket(20, 1, 1, 1) === 2)
+
+    // Test https://docs.oracle.com/cd/B28359_01/olap.111/b28126/dml_functions_2137.htm#OLADM717
+    // WIDTH_BUCKET(credit_limit, 100, 5000, 10)
+    assert(widthBucket(500, 100, 5000, 10) === 1)
+    assert(widthBucket(2300, 100, 5000, 10) === 5)
+    assert(widthBucket(3500, 100, 5000, 10) === 7)
+    assert(widthBucket(1200, 100, 5000, 10) === 3)
+    assert(widthBucket(1400, 100, 5000, 10) === 3)
+    assert(widthBucket(700, 100, 5000, 10) === 2)
+    assert(widthBucket(5000, 100, 5000, 10) === 11)
+    assert(widthBucket(1800, 100, 5000, 10) === 4)
+    assert(widthBucket(400, 100, 5000, 10) === 1)
+
+    intercept[AnalysisException]{
+      assert(widthBucket(100, 100, 5000, -1) === 1)
+    }
+  }
+}

sql/core/src/test/resources/sql-tests/inputs/operators.sql

@@ -92,3 +92,8 @@ select abs(-3.13), abs('-2.19');
 
 -- positive/negative
 select positive('-1.11'), positive(-1.11), negative('-1.11'), negative(-1.11);
+
+-- width_bucket
+select width_bucket(5.35, 0.024, 10.06, 5);
+select width_bucket(5.35, 0.024, 10.06, -5);
+select width_bucket(null, 0.024, 10.06, 5);

sql/core/src/test/resources/sql-tests/results/operators.sql.out

@@ -1,5 +1,5 @@
 -- Automatically generated by SQLQueryTestSuite
--- Number of queries: 57
+-- Number of queries: 60
 
 
 -- !query 0
@@ -468,3 +468,28 @@ select positive('-1.11'), positive(-1.11), negative('-1.11'), negative(-1.11)
 struct<(+ CAST(-1.11 AS DOUBLE)):double,(+ -1.11):decimal(3,2),(- CAST(-1.11 AS DOUBLE)):double,(- -1.11):decimal(3,2)>
 -- !query 56 output
 -1.11	-1.11	1.11	1.11
+
+
+-- !query 57
+select width_bucket(5.35, 0.024, 10.06, 5)
+-- !query 57 schema
+struct<widthbucket(CAST(5.35 AS DOUBLE), CAST(0.024 AS DOUBLE), CAST(10.06 AS DOUBLE), CAST(5 AS BIGINT)):bigint>
+-- !query 57 output
+3
+
+
+-- !query 58
+select width_bucket(5.35, 0.024, 10.06, -5)
+-- !query 58 schema
+struct<>
+-- !query 58 output
+org.apache.spark.sql.AnalysisException
+The num of bucket must be greater than 0, but got -5;
+
+
+-- !query 59
+select width_bucket(null, 0.024, 10.06, 5)
+-- !query 59 schema
+struct<widthbucket(CAST(NULL AS DOUBLE), CAST(0.024 AS DOUBLE), CAST(10.06 AS DOUBLE), CAST(5 AS BIGINT)):bigint>
+-- !query 59 output
+NULL