diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala index 3ed9ebb276a3c..d8918ef1ad2bc 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala @@ -43,7 +43,7 @@ import org.apache.spark.sql.types._ import org.apache.spark.sql.util.SQLOpenHashSet import org.apache.spark.unsafe.UTF8StringBuilder import org.apache.spark.unsafe.array.ByteArrayMethods -import org.apache.spark.unsafe.types.{ByteArray, CalendarInterval, UTF8String} +import org.apache.spark.unsafe.types.{ByteArray, CalendarInterval, TimestampNanosVal, UTF8String} /** * Base trait for [[BinaryExpression]]s with two arrays of the same element type and implicit @@ -3503,6 +3503,9 @@ case class Flatten(child: Expression) extends UnaryExpression Supported types are: byte, short, integer, long, date, timestamp. + The timestamp types include the nanosecond-precision timestamp types; their generated values + advance on the microsecond grid and keep the start value's sub-microsecond fraction. + 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' or 'year-month interval' or @@ -3580,11 +3583,8 @@ case class Sequence( val typesCorrect = DataTypeUtils.sameType(startType, stop.dataType) && (startType match { - case TimestampType | TimestampNTZType => - stepOpt.isEmpty || CalendarIntervalType.acceptsType(stepType) || - YearMonthIntervalType.acceptsType(stepType) || - DayTimeIntervalType.acceptsType(stepType) - case DateType => + case TimestampType | TimestampNTZType | DateType | + _: TimestampNTZNanosType | _: TimestampLTZNanosType => stepOpt.isEmpty || CalendarIntervalType.acceptsType(stepType) || YearMonthIntervalType.acceptsType(stepType) || DayTimeIntervalType.acceptsType(stepType) @@ -3600,7 +3600,8 @@ case class Sequence( errorSubClass = "SEQUENCE_WRONG_INPUT_TYPES", messageParameters = Map( "functionName" -> toSQLId(prettyName), - "startType" -> toSQLType(TypeCollection(TimestampType, TimestampNTZType, DateType)), + "startType" -> toSQLType( + TypeCollection(TimestampType, TimestampNTZType, AnyTimestampNanoType, DateType)), "stepType" -> toSQLType( TypeCollection(CalendarIntervalType, YearMonthIntervalType, DayTimeIntervalType)), "otherStartType" -> toSQLType(IntegralType) @@ -3630,13 +3631,21 @@ case class Sequence( val ct = physicalDataType.tag new IntegralSequenceImpl[T](iType)(ct, integral.asInstanceOf[Integral[T]]) - case TimestampType | TimestampNTZType => + case TimestampType | TimestampNTZType | + _: TimestampLTZNanosType | _: TimestampNTZNanosType => + // A nanosecond sequence reuses the microsecond machinery on epochMicros, so map each nanos + // type to its microsecond counterpart (which drives zone-aware interval addition). + val outerType: DataType = start.dataType match { + case _: TimestampLTZNanosType => TimestampType + case _: TimestampNTZNanosType => TimestampNTZType + case other => other + } if (stepOpt.isEmpty || CalendarIntervalType.acceptsType(stepOpt.get.dataType)) { - new TemporalSequenceImpl[Long](LongType, start.dataType, 1, identity, zoneId) + new TemporalSequenceImpl[Long](LongType, outerType, 1, identity, zoneId) } else if (YearMonthIntervalType.acceptsType(stepOpt.get.dataType)) { - new PeriodSequenceImpl[Long](LongType, start.dataType, 1, identity, zoneId) + new PeriodSequenceImpl[Long](LongType, outerType, 1, identity, zoneId) } else { - new DurationSequenceImpl[Long](LongType, start.dataType, 1, identity, zoneId) + new DurationSequenceImpl[Long](LongType, outerType, 1, identity, zoneId) } case DateType => @@ -3649,25 +3658,92 @@ case class Sequence( } } + private def isNanos: Boolean = start.dataType.isInstanceOf[AnyTimestampNanoType] + 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)) + if (isNanos) { + // The sequence runs on epochMicros and every element carries the start value's fraction. + // The step sign and the microsecond bound both honor the endpoints' fractions, so the result + // never overshoots stop and an out-of-order same-microsecond pair still raises the boundary + // error. See nanosStepIsNegative / nanosBoundedStopMicros. + val startNanos = startVal.asInstanceOf[TimestampNanosVal] + val stopNanos = stopVal.asInstanceOf[TimestampNanosVal] + val startMicros = startNanos.epochMicros + val startFrac = startNanos.nanosWithinMicro.toInt + // Default step sign comes from the full-precision comparison: defaultStep picks the positive + // unit when its first arg <= second, so pass (compareTo, 0) => ascending iff start <= stop. + val stepVal = stepOpt.map(_.eval(input)).getOrElse { + impl.defaultStep(startNanos.compareTo(stopNanos).toLong, 0L) + } + if (stepVal == null) return null + val stopMicros = Sequence.nanosBoundedStopMicros( + startFrac, stopNanos.epochMicros, stopNanos.nanosWithinMicro.toInt, + Sequence.nanosStepIsNegative(stepVal)) + val microsArr = impl.eval(startMicros, stopMicros, stepVal).asInstanceOf[Array[Long]] + val out = new Array[TimestampNanosVal](microsArr.length) + var i = 0 + while (i < microsArr.length) { + // startFrac is already a valid fraction, so skip the per-element range check. + out(i) = TimestampNanosVal.fromTrustedRowBytes(microsArr(i), startFrac.toShort) + i += 1 + } + ArrayData.toArrayData(out) + } else { + 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) + // Nanosecond endpoints box as TimestampNanosVal; the default-step sign uses the full-precision + // comparison (compareTo(start, stop) <= 0 => ascending), and the sequence math runs on + // epochMicros before each element is re-wrapped with the start value's fraction. + val (defaultStepStart, defaultStepStop) = if (isNanos) { + (startGen.copy(value = + JavaCode.expression(s"${startGen.value}.compareTo(${stopGen.value})", IntegerType)), + stopGen.copy(value = JavaCode.expression("0", IntegerType))) + } else { + (startGen, stopGen) + } val stepGen = stepOpt.map(_.genCode(ctx)).getOrElse( - impl.defaultStep.genCode(ctx, startGen, stopGen)) + impl.defaultStep.genCode(ctx, defaultStepStart, defaultStepStop)) val resultType = CodeGenerator.javaType(dataType) - val resultCode = { + val resultCode = if (isNanos) { + val microsArr = ctx.freshName("microsArr") + val nanosArr = ctx.freshName("nanosArr") + val startFrac = ctx.freshName("startFrac") + val startMicros = ctx.freshName("startMicros") + val stopMicros = ctx.freshName("stopMicros") + val idx = ctx.freshName("idx") + val tnv = classOf[TimestampNanosVal].getName + val genericArr = "org.apache.spark.sql.catalyst.util.GenericArrayData" + val seqObj = classOf[Sequence].getName + "$.MODULE$" + val microsGen = impl.genCode(ctx, startMicros, stopMicros, stepGen.value, microsArr, "long") + s""" + |long $startMicros = ${startGen.value}.epochMicros; + |short $startFrac = ${startGen.value}.nanosWithinMicro; + |long $stopMicros = $seqObj.nanosBoundedStopMicros( + | $startFrac, ${stopGen.value}.epochMicros, ${stopGen.value}.nanosWithinMicro, + | $seqObj.nanosStepIsNegative(${stepGen.value})); + |long[] $microsArr = null; + |$microsGen + |$tnv[] $nanosArr = new $tnv[$microsArr.length]; + |for (int $idx = 0; $idx < $microsArr.length; $idx++) { + | // startFrac is already a valid fraction, so skip the per-element range check. + | $nanosArr[$idx] = $tnv.fromTrustedRowBytes($microsArr[$idx], $startFrac); + |} + |${ev.value} = new $genericArr($nanosArr); + """.stripMargin + } else { val arr = ctx.freshName("arr") val arrElemType = CodeGenerator.javaType(dataType.elementType) s""" @@ -3746,6 +3822,30 @@ object Sequence { } } + /** + * The microsecond bound handed to the microsecond sequence machinery for a nanosecond sequence. + * Every generated element lands on the microsecond grid and carries `startFrac`, so an element + * that falls exactly on `stopMicros` is kept only when `startFrac` does not carry it past `stop` + * in the step's direction; otherwise the bound is nudged one microsecond off the boundary. That + * nudge also makes the machinery's boundary check reject an out-of-order same-microsecond pair. + */ + def nanosBoundedStopMicros( + startFrac: Int, stopMicros: Long, stopFrac: Int, stepNegative: Boolean): Long = { + if (startFrac == stopFrac) stopMicros + else if (stepNegative) if (startFrac < stopFrac) stopMicros + 1 else stopMicros + else if (startFrac > stopFrac) stopMicros - 1 else stopMicros + } + + /** Whether a sequence step points backwards, matching the microsecond machinery's sign rule. */ + def nanosStepIsNegative(step: Any): Boolean = step match { + case ci: CalendarInterval => + val totalMicros = + ci.months.toLong * (28 * MICROS_PER_DAY) + ci.days.toLong * MICROS_PER_DAY + ci.microseconds + totalMicros < 0 + case months: Int => months < 0 + case micros: Long => micros < 0 + } + private type LessThanOrEqualFn = (Any, Any) => Boolean private class DefaultStep(lteq: LessThanOrEqualFn, stepType: DataType, one: Any) { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala index 487c5233d2ce2..6fc80ca8c723b 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala @@ -37,7 +37,7 @@ import org.apache.spark.sql.errors.DataTypeErrorsBase import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.unsafe.array.ByteArrayMethods -import org.apache.spark.unsafe.types.UTF8String +import org.apache.spark.unsafe.types.{TimestampNanosVal, UTF8String} class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper with DataTypeErrorsBase { @@ -1401,6 +1401,125 @@ class CollectionExpressionsSuite Timestamp.valueOf("2018-01-01 00:00:00.000"))) } + test("SPARK-57834: sequence of nanosecond-precision timestamps") { + // The sequence advances on the microsecond grid and every element carries the start value's + // sub-microsecond fraction; the endpoints' own fractions decide the bound in full precision. + val sec = 1000000L // microseconds per second + val hour = 3600 * sec + val day = 86400000000L // microseconds per day + val month = 2678400000000L // microseconds in January 1970 (31 days), for the month-step case + def ntz(micros: Long, frac: Int, p: Int = 9): Literal = + Literal(TimestampNanosVal.fromParts(micros, frac.toShort), TimestampNTZNanosType(p)) + def ltz(micros: Long, frac: Int, p: Int = 9): Literal = + Literal(TimestampNanosVal.fromParts(micros, frac.toShort), TimestampLTZNanosType(p)) + def tnv(micros: Long, frac: Int): TimestampNanosVal = + TimestampNanosVal.fromParts(micros, frac.toShort) + + // NTZ(9), calendar-interval step; the 123ns fraction is kept on every element. + checkEvaluation(new Sequence( + ntz(0, 123), ntz(2 * sec, 123), + Literal(stringToInterval("interval 1 second"))), + Seq(tnv(0, 123), tnv(sec, 123), tnv(2 * sec, 123))) + + // Day-time interval (Duration) step and year-month interval (Period) step, exercising the + // Duration/Period sequence implementations for nanos rather than only the calendar path. + checkEvaluation(new Sequence( + ntz(0, 123), ntz(2 * hour, 123), Literal(Duration.ofHours(1))), + Seq(tnv(0, 123), tnv(hour, 123), tnv(2 * hour, 123))) + checkEvaluation(new Sequence( + ntz(0, 123), ntz(month, 123), Literal(Period.ofMonths(1))), + Seq(tnv(0, 123), tnv(month, 123))) + + // LTZ(9), hour step (zone-independent since it adds pure microseconds). + checkEvaluation(new Sequence( + ltz(0, 500), ltz(2 * hour, 500), + Literal(stringToInterval("interval 1 hour"))), + Seq(tnv(0, 500), tnv(hour, 500), tnv(2 * hour, 500))) + + // Precision 8 (fraction is a multiple of 10) and precision 7 (multiple of 100). + checkEvaluation(new Sequence( + ntz(0, 120, 8), ntz(sec, 120, 8), + Literal(stringToInterval("interval 1 second"))), + Seq(tnv(0, 120), tnv(sec, 120))) + checkEvaluation(new Sequence( + ntz(0, 100, 7), ntz(sec, 100, 7), + Literal(stringToInterval("interval 1 second"))), + Seq(tnv(0, 100), tnv(sec, 100))) + + // Differing fractions: start.frac > stop.frac drops the element that lands on stop's + // microsecond (it would exceed stop), so the result never overshoots stop. + checkEvaluation(new Sequence( + ntz(0, 900), ntz(2 * sec, 100), + Literal(stringToInterval("interval 1 second"))), + Seq(tnv(0, 900), tnv(sec, 900))) + // start.frac <= stop.frac keeps the boundary element. + checkEvaluation(new Sequence( + ntz(0, 100), ntz(2 * sec, 900), + Literal(stringToInterval("interval 1 second"))), + Seq(tnv(0, 100), tnv(sec, 100), tnv(2 * sec, 100))) + // stop on a whole second, start with a larger fraction: the element landing on stop's + // microsecond is dropped rather than overshooting stop (only start is in [start, stop]). + checkEvaluation(new Sequence( + ntz(0, 500), ntz(sec, 0), + Literal(stringToInterval("interval 1 second"))), + Seq(tnv(0, 500))) + + // Negative step, equal fractions. + checkEvaluation(new Sequence( + ntz(2 * sec, 999), ntz(0, 999), + Literal(negateExact(stringToInterval("interval 1 second")))), + Seq(tnv(2 * sec, 999), tnv(sec, 999), tnv(0, 999))) + // Negative step, differing fractions: an element on stop's microsecond is kept only when it + // stays >= stop. startFrac < stopFrac drops it; startFrac >= stopFrac keeps it. + checkEvaluation(new Sequence( + ntz(2 * sec, 100), ntz(0, 900), + Literal(negateExact(stringToInterval("interval 1 second")))), + Seq(tnv(2 * sec, 100), tnv(sec, 100))) + checkEvaluation(new Sequence( + ntz(2 * sec, 900), ntz(0, 100), + Literal(negateExact(stringToInterval("interval 1 second")))), + Seq(tnv(2 * sec, 900), tnv(sec, 900), tnv(0, 900))) + + // start == stop yields a single element that still carries the fraction. + checkEvaluation(new Sequence( + ntz(5 * sec, 42), ntz(5 * sec, 42), + Literal(stringToInterval("interval 1 second"))), + Seq(tnv(5 * sec, 42))) + + // Endpoints share a microsecond but start > stop in full precision: with a positive step this + // is an illegal boundary, matching sequence(2, 1, 1). Codegen-only, like SPARK-58440 (the + // interpreted path reports this through `require`, a plain IllegalArgumentException). The + // reported bound is nudged one microsecond off stop to trigger the machinery's boundary check. + withSQLConf( + SQLConf.CODEGEN_FACTORY_MODE.key -> CodegenObjectFactoryMode.CODEGEN_ONLY.toString) { + checkError( + exception = intercept[SparkIllegalArgumentException] { + evaluateWithMutableProjection(Sequence( + ntz(sec, 500), ntz(sec, 100), + Some(Literal(stringToInterval("interval 1 second"))), UTC_OPT)) + }, + condition = "_LEGACY_ERROR_TEMP_3243", + parameters = Map("start" -> "1000000", "stop" -> "999999", "step" -> "1000000")) + } + + // No explicit step: the default step (+1 day) is chosen from the full-precision comparison. + checkEvaluation(new Sequence(ntz(0, 7), ntz(day, 7)), + Seq(tnv(0, 7), tnv(day, 7))) + // Equal microseconds, differing fractions, no explicit step: the default step direction comes + // from the full-precision comparison, so the single in-range element is start itself (never an + // out-of-order pair). Ascending when start < stop, descending when start > stop. + checkEvaluation(new Sequence(ntz(5 * sec, 100), ntz(5 * sec, 500)), Seq(tnv(5 * sec, 100))) + checkEvaluation(new Sequence(ntz(5 * sec, 500), ntz(5 * sec, 100)), Seq(tnv(5 * sec, 500))) + + // Null propagation. + checkEvaluation(new Sequence( + Literal.create(null, TimestampNTZNanosType(9)), ntz(sec, 1), + Literal(stringToInterval("interval 1 second"))), null) + checkEvaluation(new Sequence( + ntz(0, 1), Literal.create(null, TimestampNTZNanosType(9)), + Literal(stringToInterval("interval 1 second"))), null) + } + test("Sequence on DST boundaries") { val timeZone = TimeZone.getTimeZone("Europe/Prague") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala index f20c6bd7a88ea..c29b993e500c1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala @@ -2163,6 +2163,9 @@ class DataFrameFunctionsSuite extends SharedSparkSession { Timestamp.valueOf("2018-01-02 00:00:00"))))) // test invalid data types + val nanosSeqStartType = + "(\"TIMESTAMP\" or \"TIMESTAMP_NTZ\" or " + + "\"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\" or \"DATE\")" checkError( exception = intercept[AnalysisException] { Seq((true, false)).toDF().selectExpr("sequence(_1, _2)") @@ -2171,7 +2174,7 @@ class DataFrameFunctionsSuite extends SharedSparkSession { parameters = Map( "sqlExpr" -> "\"sequence(_1, _2)\"", "functionName" -> "`sequence`", - "startType" -> "(\"TIMESTAMP\" or \"TIMESTAMP_NTZ\" or \"DATE\")", + "startType" -> nanosSeqStartType, "stepType" -> "(\"INTERVAL\" or \"INTERVAL YEAR TO MONTH\" or \"INTERVAL DAY TO SECOND\")", "otherStartType" -> "\"INTEGRAL\"" ), @@ -2185,7 +2188,7 @@ class DataFrameFunctionsSuite extends SharedSparkSession { parameters = Map( "sqlExpr" -> "\"sequence(_1, _2, _3)\"", "functionName" -> "`sequence`", - "startType" -> "(\"TIMESTAMP\" or \"TIMESTAMP_NTZ\" or \"DATE\")", + "startType" -> nanosSeqStartType, "stepType" -> "(\"INTERVAL\" or \"INTERVAL YEAR TO MONTH\" or \"INTERVAL DAY TO SECOND\")", "otherStartType" -> "\"INTEGRAL\"" ), @@ -2199,7 +2202,7 @@ class DataFrameFunctionsSuite extends SharedSparkSession { parameters = Map( "sqlExpr" -> "\"sequence(_1, _2, _3)\"", "functionName" -> "`sequence`", - "startType" -> "(\"TIMESTAMP\" or \"TIMESTAMP_NTZ\" or \"DATE\")", + "startType" -> nanosSeqStartType, "stepType" -> "(\"INTERVAL\" or \"INTERVAL YEAR TO MONTH\" or \"INTERVAL DAY TO SECOND\")", "otherStartType" -> "\"INTEGRAL\"" ), diff --git a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosFunctionsSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosFunctionsSuiteBase.scala index 4a73b2c448dfb..a997c8362a1db 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosFunctionsSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosFunctionsSuiteBase.scala @@ -1017,6 +1017,35 @@ abstract class TimestampNanosFunctionsSuiteBase extends SharedSparkSession { assert(df.selectExpr("localtimestamp() AS c").schema.head.dataType === TimestampNTZType) } } + + test("SPARK-57834: sequence over nanosecond-precision timestamps") { + val schema = new StructType() + .add("start", TimestampNTZNanosType(9)) + .add("stop", TimestampNTZNanosType(9)) + + // Every generated element keeps the start value's .000000123 fraction and the array element + // type stays TIMESTAMP_NTZ(9). + val df = spark.createDataFrame(spark.sparkContext.parallelize(Seq(Row( + LocalDateTime.parse("2020-01-01T00:00:00.000000123"), + LocalDateTime.parse("2020-01-01T00:00:02.000000123")))), schema) + val result = df.selectExpr("sequence(start, stop, INTERVAL '1' SECOND) AS s") + assert( + result.schema.head.dataType === ArrayType(TimestampNTZNanosType(9), containsNull = false)) + checkAnswer(result, Row(Seq( + LocalDateTime.parse("2020-01-01T00:00:00.000000123"), + LocalDateTime.parse("2020-01-01T00:00:01.000000123"), + LocalDateTime.parse("2020-01-01T00:00:02.000000123")))) + + // stop's smaller fraction drops the final boundary element, so the result never exceeds stop. + val df2 = spark.createDataFrame(spark.sparkContext.parallelize(Seq(Row( + LocalDateTime.parse("2020-01-01T00:00:00.000000900"), + LocalDateTime.parse("2020-01-01T00:00:02.000000100")))), schema) + checkAnswer( + df2.selectExpr("sequence(start, stop, INTERVAL '1' SECOND) AS s"), + Row(Seq( + LocalDateTime.parse("2020-01-01T00:00:00.000000900"), + LocalDateTime.parse("2020-01-01T00:00:01.000000900")))) + } } // Runs the nanosecond timestamp function tests with ANSI mode enabled explicitly.