org.apache.spark.sql.catalyst.expressions.misc.scala Maven / Gradle / Ivy
/*
* 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.expressions
import java.util.UUID
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.codegen._
import org.apache.spark.sql.catalyst.expressions.codegen.Block._
import org.apache.spark.sql.catalyst.util.RandomUUIDGenerator
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
/**
* Print the result of an expression to stderr (used for debugging codegen).
*/
case class PrintToStderr(child: Expression) extends UnaryExpression {
override def dataType: DataType = child.dataType
protected override def nullSafeEval(input: Any): Any = {
// scalastyle:off println
System.err.println(outputPrefix + input)
// scalastyle:on println
input
}
private val outputPrefix = s"Result of ${child.simpleString} is "
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
val outputPrefixField = ctx.addReferenceObj("outputPrefix", outputPrefix)
nullSafeCodeGen(ctx, ev, c =>
s"""
| System.err.println($outputPrefixField + $c);
| ${ev.value} = $c;
""".stripMargin)
}
}
/**
* A function throws an exception if 'condition' is not true.
*/
@ExpressionDescription(
usage = "_FUNC_(expr) - Throws an exception if `expr` is not true.",
examples = """
Examples:
> SELECT _FUNC_(0 < 1);
NULL
""")
case class AssertTrue(child: Expression) extends UnaryExpression with ImplicitCastInputTypes {
override def nullable: Boolean = true
override def inputTypes: Seq[DataType] = Seq(BooleanType)
override def dataType: DataType = NullType
override def prettyName: String = "assert_true"
private val errMsg = s"'${child.simpleString}' is not true!"
override def eval(input: InternalRow) : Any = {
val v = child.eval(input)
if (v == null || java.lang.Boolean.FALSE.equals(v)) {
throw new RuntimeException(errMsg)
} else {
null
}
}
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
val eval = child.genCode(ctx)
// Use unnamed reference that doesn't create a local field here to reduce the number of fields
// because errMsgField is used only when the value is null or false.
val errMsgField = ctx.addReferenceObj("errMsg", errMsg)
ExprCode(code = code"""${eval.code}
|if (${eval.isNull} || !${eval.value}) {
| throw new RuntimeException($errMsgField);
|}""".stripMargin, isNull = TrueLiteral,
value = JavaCode.defaultLiteral(dataType))
}
override def sql: String = s"assert_true(${child.sql})"
}
/**
* Returns the current database of the SessionCatalog.
*/
@ExpressionDescription(
usage = "_FUNC_() - Returns the current database.",
examples = """
Examples:
> SELECT _FUNC_();
default
""")
case class CurrentDatabase() extends LeafExpression with Unevaluable {
override def dataType: DataType = StringType
override def foldable: Boolean = true
override def nullable: Boolean = false
override def prettyName: String = "current_database"
}
// scalastyle:off line.size.limit
@ExpressionDescription(
usage = """_FUNC_() - Returns an universally unique identifier (UUID) string. The value is returned as a canonical UUID 36-character string.""",
examples = """
Examples:
> SELECT _FUNC_();
46707d92-02f4-4817-8116-a4c3b23e6266
""",
note = "The function is non-deterministic.")
// scalastyle:on line.size.limit
case class Uuid(randomSeed: Option[Long] = None) extends LeafExpression with Stateful
with ExpressionWithRandomSeed {
def this() = this(None)
override def withNewSeed(seed: Long): Uuid = Uuid(Some(seed))
override lazy val resolved: Boolean = randomSeed.isDefined
override def nullable: Boolean = false
override def dataType: DataType = StringType
@transient private[this] var randomGenerator: RandomUUIDGenerator = _
override protected def initializeInternal(partitionIndex: Int): Unit =
randomGenerator = RandomUUIDGenerator(randomSeed.get + partitionIndex)
override protected def evalInternal(input: InternalRow): Any =
randomGenerator.getNextUUIDUTF8String()
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
val randomGen = ctx.freshName("randomGen")
ctx.addMutableState("org.apache.spark.sql.catalyst.util.RandomUUIDGenerator", randomGen,
forceInline = true,
useFreshName = false)
ctx.addPartitionInitializationStatement(s"$randomGen = " +
"new org.apache.spark.sql.catalyst.util.RandomUUIDGenerator(" +
s"${randomSeed.get}L + partitionIndex);")
ev.copy(code = code"final UTF8String ${ev.value} = $randomGen.getNextUUIDUTF8String();",
isNull = FalseLiteral)
}
override def freshCopy(): Uuid = Uuid(randomSeed)
}