Skip to content

Commit c6ba7cc

Browse files
adrian-wangrxin
authored andcommitted
[SPARK-8215] [SPARK-8212] [SQL] add leaf math expression for e and pi
Author: Daoyuan Wang <[email protected]> Closes apache#6716 from adrian-wang/epi and squashes the following commits: e2e8dbd [Daoyuan Wang] move tests 11b351c [Daoyuan Wang] add tests and remove pu db331c9 [Daoyuan Wang] py style 599ddd8 [Daoyuan Wang] add py e6783ef [Daoyuan Wang] register function 82d426e [Daoyuan Wang] add function entry dbf3ab5 [Daoyuan Wang] add PI and E
1 parent e90035e commit c6ba7cc

File tree

5 files changed

+96
-0
lines changed

5 files changed

+96
-0
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,13 +106,15 @@ object FunctionRegistry {
106106
expression[Cbrt]("cbrt"),
107107
expression[Ceil]("ceil"),
108108
expression[Cos]("cos"),
109+
expression[EulerNumber]("e"),
109110
expression[Exp]("exp"),
110111
expression[Expm1]("expm1"),
111112
expression[Floor]("floor"),
112113
expression[Hypot]("hypot"),
113114
expression[Log]("log"),
114115
expression[Log10]("log10"),
115116
expression[Log1p]("log1p"),
117+
expression[Pi]("pi"),
116118
expression[Pow]("pow"),
117119
expression[Rint]("rint"),
118120
expression[Signum]("signum"),

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,34 @@ package org.apache.spark.sql.catalyst.expressions
2020
import org.apache.spark.sql.catalyst.expressions.codegen._
2121
import org.apache.spark.sql.types.{DataType, DoubleType}
2222

23+
/**
24+
* A leaf expression specifically for math constants. Math constants expect no input.
25+
* @param c The math constant.
26+
* @param name The short name of the function
27+
*/
28+
abstract class LeafMathExpression(c: Double, name: String)
29+
extends LeafExpression with Serializable {
30+
self: Product =>
31+
32+
override def dataType: DataType = DoubleType
33+
override def foldable: Boolean = true
34+
override def nullable: Boolean = false
35+
override def toString: String = s"$name()"
36+
37+
override def eval(input: Row): Any = c
38+
39+
override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
40+
s"""
41+
boolean ${ev.isNull} = false;
42+
${ctx.javaType(dataType)} ${ev.primitive} = java.lang.Math.$name;
43+
"""
44+
}
45+
}
46+
2347
/**
2448
* A unary expression specifically for math functions. Math Functions expect a specific type of
2549
* input format, therefore these functions extend `ExpectsInputTypes`.
50+
* @param f The math function.
2651
* @param name The short name of the function
2752
*/
2853
abstract class UnaryMathExpression(f: Double => Double, name: String)
@@ -98,6 +123,16 @@ abstract class BinaryMathExpression(f: (Double, Double) => Double, name: String)
98123
}
99124
}
100125

126+
////////////////////////////////////////////////////////////////////////////////////////////////////
127+
////////////////////////////////////////////////////////////////////////////////////////////////////
128+
// Leaf math functions
129+
////////////////////////////////////////////////////////////////////////////////////////////////////
130+
////////////////////////////////////////////////////////////////////////////////////////////////////
131+
132+
case class EulerNumber() extends LeafMathExpression(math.E, "E")
133+
134+
case class Pi() extends LeafMathExpression(math.Pi, "PI")
135+
101136
////////////////////////////////////////////////////////////////////////////////////////////////////
102137
////////////////////////////////////////////////////////////////////////////////////////////////////
103138
// Unary math functions

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,20 @@ import org.apache.spark.sql.types.DoubleType
2222

2323
class MathFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
2424

25+
/**
26+
* Used for testing leaf math expressions.
27+
*
28+
* @param e expression
29+
* @param c The constants in scala.math
30+
* @tparam T Generic type for primitives
31+
*/
32+
private def testLeaf[T](
33+
e: () => Expression,
34+
c: T): Unit = {
35+
checkEvaluation(e(), c, EmptyRow)
36+
checkEvaluation(e(), c, create_row(null))
37+
}
38+
2539
/**
2640
* Used for testing unary math expressions.
2741
*
@@ -74,6 +88,14 @@ class MathFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
7488
checkEvaluation(c(Literal(1.0), Literal.create(null, DoubleType)), null, create_row(null))
7589
}
7690

91+
test("e") {
92+
testLeaf(EulerNumber, math.E)
93+
}
94+
95+
test("pi") {
96+
testLeaf(Pi, math.Pi)
97+
}
98+
7799
test("sin") {
78100
testUnary(Sin, math.sin)
79101
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -944,6 +944,15 @@ object functions {
944944
*/
945945
def cosh(columnName: String): Column = cosh(Column(columnName))
946946

947+
/**
948+
* Returns the double value that is closer than any other to e, the base of the natural
949+
* logarithms.
950+
*
951+
* @group math_funcs
952+
* @since 1.5.0
953+
*/
954+
def e(): Column = EulerNumber()
955+
947956
/**
948957
* Computes the exponential of the given value.
949958
*
@@ -1105,6 +1114,15 @@ object functions {
11051114
*/
11061115
def log1p(columnName: String): Column = log1p(Column(columnName))
11071116

1117+
/**
1118+
* Returns the double value that is closer than any other to pi, the ratio of the circumference
1119+
* of a circle to its diameter.
1120+
*
1121+
* @group math_funcs
1122+
* @since 1.5.0
1123+
*/
1124+
def pi(): Column = Pi()
1125+
11081126
/**
11091127
* Returns the value of the first argument raised to the power of the second argument.
11101128
*

sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,25 @@ class DataFrameFunctionsSuite extends QueryTest {
8585
}
8686
}
8787

88+
test("constant functions") {
89+
checkAnswer(
90+
testData2.select(e()).limit(1),
91+
Row(scala.math.E)
92+
)
93+
checkAnswer(
94+
testData2.select(pi()).limit(1),
95+
Row(scala.math.Pi)
96+
)
97+
checkAnswer(
98+
ctx.sql("SELECT E()"),
99+
Row(scala.math.E)
100+
)
101+
checkAnswer(
102+
ctx.sql("SELECT PI()"),
103+
Row(scala.math.Pi)
104+
)
105+
}
106+
88107
test("bitwiseNOT") {
89108
checkAnswer(
90109
testData2.select(bitwiseNOT($"a")),

0 commit comments

Comments
 (0)