Merge pull request #11806 from NVIDIA/branch-24.12

[auto-merge] branch-24.12 to branch-25.02 [skip ci] [bot]

integration_tests/src/main/python/date_time_test.py

@@ -139,34 +139,34 @@ def test_datediff(data_gen):
 
 hms_fallback = ['ProjectExec'] if not is_supported_time_zone() else []
 
-@allow_non_gpu(*hms_fallback)
+@allow_non_gpu(*non_utc_tz_allow)
 def test_months_between():
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : binary_op_df(spark, timestamp_gen).selectExpr('months_between(a, b, false)'))
 
-@allow_non_gpu(*hms_fallback)
+@allow_non_gpu(*non_utc_tz_allow)
 def test_months_between_first_day():
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('months_between(a, timestamp"2024-01-01", false)'))
 
-@allow_non_gpu(*hms_fallback)
+@allow_non_gpu(*non_utc_tz_allow)
 def test_months_between_last_day():
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('months_between(a, timestamp"2023-12-31", false)'))
 
-@allow_non_gpu(*hms_fallback)
+@allow_non_gpu(*non_utc_tz_allow)
 @approximate_float()
 def test_months_between_round():
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : binary_op_df(spark, timestamp_gen).selectExpr('months_between(a, b, true)'))
 
-@allow_non_gpu(*hms_fallback)
+@allow_non_gpu(*non_utc_tz_allow)
 @approximate_float()
 def test_months_between_first_day_round():
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('months_between(a, timestamp"2024-01-01", true)'))
 
-@allow_non_gpu(*hms_fallback)
+@allow_non_gpu(*non_utc_tz_allow)
 @approximate_float()
 def test_months_between_last_day_round():
     assert_gpu_and_cpu_are_equal_collect(

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala

@@ -1217,7 +1217,8 @@ class MonthsBetweenExprMeta(expr: MonthsBetween,
                             rule: DataFromReplacementRule)
   extends ExprMeta[MonthsBetween](expr, conf, parent, rule) {
 
-  override def isTimeZoneSupported = true
+  // See https://github.com/NVIDIA/spark-rapids/issues/11800
+  override def isTimeZoneSupported = false
 
   override def convertToGpu(): GpuExpression = {
     val gpuChildren = childExprs.map(_.convertToGpu())
@@ -1287,6 +1288,13 @@ object GpuMonthsBetween {
   private def calcSecondsInDay(converted: ColumnVector): ColumnVector = {
     // Find the number of seconds that are not counted for in a day
 
+    // Rounding down to the current day, only works if you are in a time zone with no
+    // transition rules. This is because if a transition happens in between the start
+    // of the day and the timestamp we will be off. As such this will need to change to
+    // support other time zones, and it will need to take the timezone into account when
+    // calculating this.
+    // https://github.com/NVIDIA/spark-rapids/issues/11800
+
     // find the micros over by finding the part that is not days
     val microsInDay = withResource(converted.dateTimeFloor(DateTimeRoundingFrequency.DAY)) { days =>
       // But we cannot subtract timestamps directly. They are both micros