Fix UNBOUNDED time-range query. (NVIDIA#1169)

Signed-off-by: Mithun RK <mythrocks@gmail.com>
nartal1 · Dec 3, 2020 · 3f0d1d1 · 3f0d1d1
1 parent f799679
commit 3f0d1d1
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 15 deletions.
diff --git a/integration_tests/src/main/python/window_function_test.py b/integration_tests/src/main/python/window_function_test.py
@@ -144,12 +144,13 @@ def do_it(spark):
                 .withColumn('row_num', f.row_number().over(baseWindowSpec))
     assert_gpu_and_cpu_are_equal_collect(do_it, conf={'spark.rapids.sql.hasNans': 'false'})
 
+
 # Test for RANGE queries, with timestamp order-by expressions.
 # Non-timestamp order-by columns are currently unsupported for RANGE queries.
 # See https://github.com/NVIDIA/spark-rapids/issues/216
 @ignore_order
 @pytest.mark.parametrize('data_gen', [_grpkey_longs_with_timestamps,
-                                      pytest.param(_grpkey_longs_with_nullable_timestamps, marks=pytest.mark.xfail(reason='https://github.com/NVIDIA/spark-rapids/issues/1039'))],
+                                      pytest.param(_grpkey_longs_with_nullable_timestamps)],
                                       ids=idfn)
 def test_window_aggs_for_ranges(data_gen):
     assert_gpu_and_cpu_are_equal_sql(

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWindowExpression.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWindowExpression.scala
@@ -135,17 +135,6 @@ class GpuWindowExpressionMeta(
                   willNotWorkOnGpu("only a single date/time based column in window" +
                       " range functions is supported")
                 }
-
-                // https://github.com/NVIDIA/spark-rapids/issues/1039
-                // The order by column does not work for nullable time range queries that include
-                // UNBOUNDED
-                // Once this is fixed please uncomment the tests in WindowFunctionSuite that
-                // are impacted by this change
-                if (orderSpec.exists(_.nullable) &&
-                    (lower == Int.MinValue || upper == Int.MaxValue)) {
-                  willNotWorkOnGpu("nullable date/timestamp ranges" +
-                      " are not supported with UNBOUNDED bounds")
-                }
               } else {
                 willNotWorkOnGpu("a mixture of date/time and non date/time based" +
                     " columns is not supported in a window range function")
@@ -357,9 +346,23 @@ object GpuWindowExpression {
     val lower = getRangeBasedLower(windowFrameSpec)
     val upper = getRangeBasedUpper(windowFrameSpec)
 
-    val windowOptionBuilder = WindowOptions.builder().minPeriods(1)
-        .window(lower, upper)
-        .timestampColumnIndex(timeColumnIndex)
+    val windowOptionBuilder = WindowOptions.builder()
+                                .minPeriods(1)
+                                .timestampColumnIndex(timeColumnIndex)
+
+    if (lower.equals(Int.MaxValue)) {
+      windowOptionBuilder.unboundedPreceding()
+    }
+    else {
+      windowOptionBuilder.preceding(lower)
+    }
+
+    if (upper.equals(Int.MaxValue)) {
+      windowOptionBuilder.unboundedFollowing()
+    }
+    else {
+      windowOptionBuilder.following(upper)
+    }
 
     // We only support a single time based column to order by right now, so just verify
     // that it is correct.