Let Logarithm accept one parameter too.

apache · Jun 11, 2015 · ebc9929 · ebc9929
1 parent 605574d
commit ebc9929
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 4 deletions.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/math.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/math.scala
@@ -261,11 +261,47 @@ object Logarithm {
 
 case class Logarithm(left: Expression, right: Expression)
   extends BinaryMathExpression((c1, c2) => math.log(c2) / math.log(c1), "LOG") {
+  override def eval(input: Row): Any = {
+    val evalE2 = right.eval(input)
+    if (evalE2 == null) {
+      null
+    } else {
+      val evalE1 = left.eval(input)
+      var result: Double = 0.0
+      if (evalE1 == null) {
+        result = math.log(evalE2.asInstanceOf[Double])
+      } else {
+        result = math.log(evalE2.asInstanceOf[Double]) / math.log(evalE1.asInstanceOf[Double])
+      }
+      if (result.isNaN) null else result
+    }
+  }
+
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    defineCodeGen(ctx, ev, (c1, c2) => s"java.lang.Math.log($c2) / java.lang.Math.log($c1)") + s"""
+    if (left.dataType != right.dataType) {
+      // log.warn(s"${left.dataType} != ${right.dataType}")
+    }
+
+    val eval1 = left.gen(ctx)
+    val eval2 = right.gen(ctx)
+    val resultCode =
+      s"java.lang.Math.log(${eval2.primitive}) / java.lang.Math.log(${eval1.primitive})"
+
+    s"""
+      ${eval2.code}
+      boolean ${ev.isNull} = ${eval2.isNull};
+      ${ctx.javaType(dataType)} ${ev.primitive} = ${ctx.defaultValue(dataType)};
+      if (!${ev.isNull}) {
+        ${eval1.code}
+        if (!${eval1.isNull}) {
+          ${ev.primitive} = ${resultCode};
+        } else {
+          ${ev.primitive} = java.lang.Math.log(${eval2.primitive});
+        }
+      }
       if (Double.valueOf(${ev.primitive}).isNaN()) {
         ${ev.isNull} = true;
       }
-      """
+    """
   }
 }
diff --git a/...atalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathFunctionsSuite.scala b/...atalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathFunctionsSuite.scala
@@ -205,7 +205,17 @@ class MathFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("binary log") {
-    testBinary((e1, e2) => Logarithm(e1, e2), (c1, c2) => math.log(c2) / math.log(c1),
-      (1 to 20).map(v => (v * 0.1, v * 0.2)))
+    val f =  (c1: Double, c2: Double) => math.log(c2) / math.log(c1)
+    val domain = (1 to 20).map(v => (v * 0.1, v * 0.2))
+
+    domain.foreach { case (v1, v2) =>
+        checkEvaluation(Logarithm(Literal(v1), Literal(v2)), f(v1 + 0.0, v2 + 0.0), EmptyRow)
+        checkEvaluation(Logarithm(Literal(v2), Literal(v1)), f(v2 + 0.0, v1 + 0.0), EmptyRow)
+    }
+    // When base is null, Logarithm is as same as Log
+    checkEvaluation(Logarithm(Literal.create(null, DoubleType), Literal(1.0)),
+      math.log(1.0), create_row(null))
+    checkEvaluation(Logarithm(Literal(1.0), Literal.create(null, DoubleType)),
+      null, create_row(null))
   }
 }