From 43175912a9ef5813c884f3ea4f689ee8b8d22b50 Mon Sep 17 00:00:00 2001 From: "Robert (Bobby) Evans" Date: Tue, 16 Feb 2021 16:19:09 -0600 Subject: [PATCH] Disable range partitioning if computation is needed (#1734) Signed-off-by: Robert (Bobby) Evans --- integration_tests/src/main/python/sort_test.py | 12 ++++++++++++ .../scala/com/nvidia/spark/rapids/GpuOverrides.scala | 12 ++++++++++++ 2 files changed, 24 insertions(+) diff --git a/integration_tests/src/main/python/sort_test.py b/integration_tests/src/main/python/sort_test.py index 7c1ebd2a81e..9f9557f2eed 100644 --- a/integration_tests/src/main/python/sort_test.py +++ b/integration_tests/src/main/python/sort_test.py @@ -67,3 +67,15 @@ def test_multi_orderby(data_gen): def test_multi_orderby_with_limit(data_gen): assert_gpu_and_cpu_are_equal_collect( lambda spark : binary_op_df(spark, data_gen).orderBy(f.col('a'), f.col('b').desc()).limit(100)) + +# We are not trying all possibilities, just doing a few with numbers so the query works. +@pytest.mark.parametrize('data_gen', [byte_gen, long_gen, float_gen], ids=idfn) +def test_orderby_with_processing(data_gen): + assert_gpu_and_cpu_are_equal_collect( + lambda spark : unary_op_df(spark, data_gen).orderBy(f.lit(100) - f.col('a'))) + +# We are not trying all possibilities, just doing a few with numbers so the query works. +@pytest.mark.parametrize('data_gen', [byte_gen, long_gen, float_gen], ids=idfn) +def test_orderby_with_processing_and_limit(data_gen): + assert_gpu_and_cpu_are_equal_collect( + lambda spark : unary_op_df(spark, data_gen).orderBy(f.lit(100) - f.col('a')).limit(100)) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala index 037cd4b7440..903ab71230b 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala @@ -2386,6 +2386,18 @@ object GpuOverrides { override val childExprs: Seq[BaseExprMeta[_]] = rp.ordering.map(GpuOverrides.wrapExpr(_, conf, Some(this))) + override def tagPartForGpu(): Unit = { + def isSortOrderSimpleEnough(so: SortOrder): Boolean = so.child match { + case _: AttributeReference => true + case _ => false + } + // Once https://github.com/NVIDIA/spark-rapids/issues/1730 is fixed this check should be + // removed + if (!rp.ordering.forall(isSortOrderSimpleEnough)) { + willNotWorkOnGpu("computation is not supported for sort order in range partitioning") + } + } + override def convertToGpu(): GpuPartitioning = { if (rp.numPartitions > 1) { val gpuOrdering = childExprs.map(_.convertToGpu()).asInstanceOf[Seq[SortOrder]]