Disable range partitioning if computation is needed (NVIDIA#1734)

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

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))

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]]