diff --git a/integration_tests/src/main/python/window_function_test.py b/integration_tests/src/main/python/window_function_test.py index 1c3f38fd793..1c5b170dfa9 100644 --- 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 index a4bc79deeb0..e0255637a9d 100644 --- 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.