Github user ueshin commented on a diff in the pull request:
https://github.com/apache/spark/pull/21155#discussion_r193927263
--- Diff:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
---
@@ -1887,6 +1889,402 @@ case class Flatten(child: Expression) extends
UnaryExpression {
override def prettyName: String = "flatten"
}
+@ExpressionDescription(
+ usage = """
+ _FUNC_(start, stop, step) - Generates an array of elements from start
to stop (inclusive),
+ incrementing by step. The type of the returned elements is the same
as the type of argument
+ expressions.
+
+ Supported types are: byte, short, integer, long, date, timestamp.
+
+ The start and stop expressions must resolve to the same type.
+ If start and stop expressions resolve to the 'date' or 'timestamp'
type
+ then the step expression must resolve to the 'interval' type,
otherwise to the same type
+ as the start and stop expressions.
+ """,
+ arguments = """
+ Arguments:
+ * start - an expression. The start of the range.
+ * stop - an expression. The end the range (inclusive).
+ * step - an optional expression. The step of the range.
+ By default step is 1 if start is less than or equal to stop,
otherwise -1.
+ For the temporal sequences it's 1 day and -1 day respectively.
+ If start is greater than stop then the step must be negative,
and vice versa.
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(1, 5);
+ [1, 2, 3, 4, 5]
+ > SELECT _FUNC_(5, 1);
+ [5, 4, 3, 2, 1]
+ > SELECT _FUNC_(to_date('2018-01-01'), to_date('2018-03-01'),
interval 1 month);
+ [2018-01-01, 2018-02-01, 2018-03-01]
+ """,
+ since = "2.4.0"
+)
+case class Sequence(
+ start: Expression,
+ stop: Expression,
+ stepOpt: Option[Expression],
+ timeZoneId: Option[String] = None)
+ extends Expression
+ with TimeZoneAwareExpression {
+
+ import Sequence._
+
+ def this(start: Expression, stop: Expression) =
+ this(start, stop, None, None)
+
+ def this(start: Expression, stop: Expression, step: Expression) =
+ this(start, stop, Some(step), None)
+
+ override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression =
+ copy(timeZoneId = Some(timeZoneId))
+
+ override def children: Seq[Expression] = Seq(start, stop) ++ stepOpt
+
+ override def foldable: Boolean = children.forall(_.foldable)
+
+ override def nullable: Boolean = children.exists(_.nullable)
+
+ override lazy val dataType: ArrayType = ArrayType(start.dataType,
containsNull = false)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val startType = start.dataType
+ def stepType = stepOpt.get.dataType
+ val typesCorrect =
+ startType.sameType(stop.dataType) &&
+ (startType match {
+ case TimestampType | DateType =>
+ stepOpt.isEmpty || CalendarIntervalType.acceptsType(stepType)
+ case _: IntegralType =>
+ stepOpt.isEmpty || stepType.sameType(startType)
+ case _ => false
+ })
+
+ if (typesCorrect) {
+ TypeCheckResult.TypeCheckSuccess
+ } else {
+ TypeCheckResult.TypeCheckFailure(
+ s"$prettyName only supports integral, timestamp or date types")
+ }
+ }
+
+ def coercibleChildren: Seq[Expression] = children.filter(_.dataType !=
CalendarIntervalType)
+
+ def castChildrenTo(widerType: DataType): Expression = Sequence(
+ Cast(start, widerType),
+ Cast(stop, widerType),
+ stepOpt.map(step => if (step.dataType != CalendarIntervalType)
Cast(step, widerType) else step),
+ timeZoneId)
+
+ private lazy val impl: SequenceImpl = dataType.elementType match {
+ case iType: IntegralType =>
+ type T = iType.InternalType
+ val ct = ClassTag[T](iType.tag.mirror.runtimeClass(iType.tag.tpe))
+ new IntegralSequenceImpl(iType)(ct, iType.integral)
+
+ case TimestampType =>
+ new TemporalSequenceImpl[Long](LongType, 1, identity, timeZone)
+
+ case DateType =>
+ new TemporalSequenceImpl[Int](IntegerType, MICROS_PER_DAY, _.toInt,
timeZone)
+ }
+
+ override def eval(input: InternalRow): Any = {
+ val startVal = start.eval(input)
+ if (startVal == null) return null
+ val stopVal = stop.eval(input)
+ if (stopVal == null) return null
+ val stepVal =
stepOpt.map(_.eval(input)).getOrElse(impl.defaultStep(startVal, stopVal))
+ if (stepVal == null) return null
+
+ ArrayData.toArrayData(impl.eval(startVal, stopVal, stepVal))
+ }
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ val startGen = start.genCode(ctx)
+ val stopGen = stop.genCode(ctx)
+ val stepGen = stepOpt.map(_.genCode(ctx)).getOrElse(
+ impl.defaultStep.genCode(ctx, startGen, stopGen))
+
+ val resultType = CodeGenerator.javaType(dataType)
+ val resultCode = {
+ val arr = ctx.freshName("arr")
+ val arrElemType = CodeGenerator.javaType(dataType.elementType)
+ s"""
+ |final $arrElemType[] $arr = null;
+ |${impl.genCode(ctx, startGen.value, stopGen.value,
stepGen.value, arr, arrElemType)}
+ |${ev.value} = UnsafeArrayData.fromPrimitiveArray($arr);
+ """.stripMargin
+ }
+
+ if (nullable) {
+ val nullSafeEval =
+ startGen.code + ctx.nullSafeExec(start.nullable, startGen.isNull) {
+ stopGen.code + ctx.nullSafeExec(stop.nullable, stopGen.isNull) {
+ stepGen.code + ctx.nullSafeExec(stepOpt.exists(_.nullable),
stepGen.isNull) {
+ s"""
+ |${ev.isNull} = false;
+ |$resultCode
+ """.stripMargin
+ }
+ }
+ }
+ ev.copy(code =
+ s"""
+ |boolean ${ev.isNull} = true;
+ |$resultType ${ev.value} = null;
+ |$nullSafeEval
+ """.stripMargin)
+
+ } else {
+ ev.copy(code =
+ s"""
+ |boolean ${ev.isNull} = false;
--- End diff --
Maybe we don't need this?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
