stevomitric commented on code in PR #56811: URL: https://github.com/apache/spark/pull/56811#discussion_r3491399997
########## sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosWindowSuiteBase.scala: ########## @@ -0,0 +1,313 @@ +/* + * 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 + +import java.time.{Instant, LocalDateTime} + +import org.apache.spark.SparkConf +import org.apache.spark.sql.expressions.Window +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ + +/** + * End-to-end window-function correctness tests over the nanosecond-precision timestamp types + * `TIMESTAMP_NTZ(p)` / `TIMESTAMP_LTZ(p)` (`p` in `[7, 9]`). Window functions are type-agnostic + * and ride entirely on orderability and the `UnsafeRow` / window-buffer primitives, so no + * production change is required -- this suite locks the behaviour in. + * + * The headline assertion is sub-microsecond ordering: input values share their `epochMicros` and + * differ only in `nanosWithinMicro`, so the micro path cannot distinguish them and + * `row_number()` / `rank()` / `dense_rank()` are the real proof of nanos ordering. + * `lag` / `lead` additionally round-trip the nanos value through the window buffer / `UnsafeRow` + * append, so collecting the neighbour back as `LocalDateTime` / `Instant` proves the carrier + * (`epochMicros` + `nanosWithinMicro`) survives. Each ordering body runs on both the whole-stage + * codegen comparison arm and the interpreted `Ordering[TimestampNanosVal]` arm, NTZ and LTZ. + * + * All sub-microsecond remainders are multiples of 100ns (100 / 200 / ... / 900) so they are exact + * at every precision p in [7, 9] (p=7 has 100ns resolution, p=8 has 10ns); a non-100ns-multiple + * remainder would be floored away at p=7/p=8 and collapse the intended distinct values into ties. + * + * The preview flag is enabled by default under tests (`Utils.isTesting`), so it is not set. The + * session time zone is fixed so `TIMESTAMP_LTZ` values render deterministically. The two + * subclasses run every test with ANSI mode on and off. + * + * NOTE: every test here projects a deterministic, distinct ordinal column (`id`, or the window + * output `rn`/`rk`) alongside the nanos column, so `checkAnswer` (order-insensitive) suffices -- + * the row-number / rank value IS the ordering proof, so no collect-strict assertion is needed. + */ +abstract class TimestampNanosWindowSuiteBase extends SharedSparkSession { + + import testImplicits._ + + override def sparkConf: SparkConf = super.sparkConf + .set(SQLConf.SESSION_LOCAL_TIMEZONE.key, "America/Los_Angeles") + + protected val codegenModes: Seq[Seq[(String, String)]] = Seq( + Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY"), + Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN")) + + private def ntzSchema(p: Int): StructType = + new StructType().add("id", IntegerType).add("ts", TimestampNTZNanosType(p)) + + private def ltzSchema(p: Int): StructType = + new StructType().add("id", IntegerType).add("ts", TimestampLTZNanosType(p)) + + // ========================================================================================== + // row_number() OVER (ORDER BY <nanos col>) -- sub-microsecond ordering, NTZ + LTZ. + // ========================================================================================== + // All three values share epochMicros 2020-01-01T00:00:00.000000 and differ only inside the + // microsecond (100ns / 500ns / 900ns), so the row numbers are produced purely by nanos ordering. + test("row_number over a nanosecond TIMESTAMP_NTZ orders by the sub-microsecond part") { + codegenModes.foreach { conf => + withSQLConf(conf: _*) { + Seq(7, 8, 9).foreach { p => + val data = Seq( + Row(10, LocalDateTime.parse("2020-01-01T00:00:00.000000900")), + Row(20, LocalDateTime.parse("2020-01-01T00:00:00.000000100")), + Row(30, LocalDateTime.parse("2020-01-01T00:00:00.000000500"))) + val df = spark.createDataFrame(spark.sparkContext.parallelize(data), ntzSchema(p)) + // ASC: 100ns -> 500ns -> 900ns -> ids 20, 30, 10. + checkAnswer( + df.select($"id", row_number().over(Window.orderBy($"ts")).as("rn")), + Seq(Row(20, 1), Row(30, 2), Row(10, 3))) + // DESC: reversed. + checkAnswer( + df.select($"id", row_number().over(Window.orderBy($"ts".desc)).as("rn")), + Seq(Row(10, 1), Row(30, 2), Row(20, 3))) + } + } + } + } + + test("row_number over a nanosecond TIMESTAMP_LTZ orders by the sub-microsecond part") { Review Comment: added the DESC arm to the LTZ row_number test. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
