Support group-limit optimization for ROW_NUMBER (#11886)

* Support group-limit optimization for `ROW_NUMBER`

Fixes #10505.

This is a follow-up to #10500, which added support for WindowGroupLimit optimizations for `RANK` and `DENSE_RANK` window functions.  The same optimization was not extended to `ROW_NUMBER` at that time.

This commit now allows the output from `ROW_NUMBER` to be filtered map-side, in case there is a `<` predicate on its return value.
The following is an example of the kind of query that is affected by this change:
```sql
SELECT * FROM (
  SELECT *,
         ROW_NUMBER() OVER (PARTITION BY p ORDER BY o) AS rn
  FROM mytable )
WHERE rn < 10;
```

This is per the optimization in [SPARK-37099](https://issues.apache.org/jira/browse/SPARK-37099) in Apache Spark.  With this, the output from the window function could potentially be drastically smaller than the input, thus saving on shuffle traffic.

Note that this optimization does not kick in on Apache Spark or in `spark-rapids` if the `ROW_NUMBER` phrase does not include a `PARTITION BY` clause.  `spark-rapids` remains consistent with Apache Spark in this regard.

Signed-off-by: MithunR <[email protected]>

integration_tests/src/main/python/window_function_test.py

@@ -2107,7 +2107,9 @@ def assert_query_runs_on(exec, conf):
                             'DENSE_RANK() OVER (PARTITION BY a ORDER BY b, c) ',
                             'RANK() OVER (ORDER BY a,b,c) ',
                             'DENSE_RANK() OVER (ORDER BY a,b,c) ',
-                        ])
+                            # ROW_NUMBER() on an un-partitioned window does not invoke WindowGroupLimit optimization.
+                            'ROW_NUMBER() OVER (PARTITION BY a ORDER BY b,c) ',
+])
 def test_window_group_limits_for_ranking_functions(data_gen, batch_size, rank_clause):
     """
     This test verifies that window group limits are applied for queries with ranking-function based
@@ -2133,39 +2135,6 @@ def test_window_group_limits_for_ranking_functions(data_gen, batch_size, rank_cl
         conf=conf)
 
 
-@allow_non_gpu('WindowGroupLimitExec')
-@pytest.mark.skipif(condition=not (is_spark_350_or_later() or is_databricks133_or_later()),
-                    reason="WindowGroupLimit not available for spark.version < 3.5 "
-                           " and Databricks version < 13.3")
-@ignore_order(local=True)
-@approximate_float
-def test_window_group_limits_fallback_for_row_number():
-    """
-    This test verifies that window group limits are applied for queries with ranking-function based
-    row filters.
-    This test covers RANK() and DENSE_RANK(), for window function with and without `PARTITIONED BY`
-    clauses.
-    """
-    conf = {'spark.rapids.sql.batchSizeBytes': '1g',
-            'spark.rapids.sql.castFloatToDecimal.enabled': True}
-
-    data_gen = _grpkey_longs_with_no_nulls
-    query = """
-        SELECT * FROM (
-          SELECT *, ROW_NUMBER() OVER (PARTITION BY a ORDER BY b, c) AS rnk
-          FROM window_agg_table
-        )
-        WHERE rnk < 3
-     """
-
-    assert_gpu_sql_fallback_collect(
-        lambda spark: gen_df(spark, data_gen, length=512),
-        cpu_fallback_class_name="WindowGroupLimitExec",
-        table_name="window_agg_table",
-        sql=query,
-        conf=conf)
-
-
 def test_lru_cache_datagen():
     # log cache info at the end of integration tests, not related to window functions
     info = gen_df_help.cache_info()

sql-plugin/src/main/spark341db/scala/com/nvidia/spark/rapids/shims/GpuWindowGroupLimitExec.scala

@@ -30,12 +30,12 @@ import ai.rapids.cudf.GroupByOptions
 import com.nvidia.spark.rapids.{BaseExprMeta, DataFromReplacementRule, GpuBindReferences, GpuBoundReference, GpuColumnVector, GpuExec, GpuExpression, GpuMetric, GpuOverrides, GpuProjectExec, RapidsConf, RapidsMeta, SparkPlanMeta, SpillableColumnarBatch, SpillPriorities}
 import com.nvidia.spark.rapids.Arm.withResource
 import com.nvidia.spark.rapids.RmmRapidsRetryIterator.{splitSpillableInHalfByRows, withRetry}
-import com.nvidia.spark.rapids.window.{GpuDenseRank, GpuRank}
+import com.nvidia.spark.rapids.window.{GpuDenseRank, GpuRank, GpuRowNumber}
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Attribute, DenseRank, Expression, Rank, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, DenseRank, Expression, Rank, RowNumber, SortOrder}
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.window.{Final, Partial, WindowGroupLimitExec, WindowGroupLimitMode}
 import org.apache.spark.sql.types.DataType
@@ -44,6 +44,7 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
 sealed trait RankFunctionType
 case object RankFunction extends RankFunctionType
 case object DenseRankFunction extends RankFunctionType
+case object RowNumberFunction extends RankFunctionType
 
 class GpuWindowGroupLimitExecMeta(limitExec: WindowGroupLimitExec,
                                   conf: RapidsConf,
@@ -63,8 +64,8 @@ class GpuWindowGroupLimitExecMeta(limitExec: WindowGroupLimitExec,
     wrapped.rankLikeFunction match {
       case DenseRank(_) =>
       case Rank(_) =>
-      // case RowNumber() => // TODO: Future.
-      case _ => willNotWorkOnGpu("Only Rank() and DenseRank() are " +
+      case RowNumber() =>
+      case _ => willNotWorkOnGpu("Only rank, dense_rank and row_number are " +
                                  "currently supported for window group limits")
     }
 
@@ -166,6 +167,7 @@ class GpuWindowGroupLimitingIterator(input: Iterator[ColumnarBatch],
     rankFunctionType match {
       case RankFunction => GpuRank(sortColumns)
       case DenseRankFunction => GpuDenseRank(sortColumns)
+      case RowNumberFunction => GpuRowNumber
       case _ => throw new IllegalArgumentException("Unexpected ranking function")
     }
   }
@@ -300,8 +302,10 @@ case class GpuWindowGroupLimitExec(
   private def getRankFunctionType(expr: Expression): RankFunctionType = expr match {
     case GpuRank(_) => RankFunction
     case GpuDenseRank(_) => DenseRankFunction
+    case GpuRowNumber => RowNumberFunction
     case _ =>
-      throw new UnsupportedOperationException("Only Rank() is currently supported for group limits")
+      throw new UnsupportedOperationException("Only rank, dense_rank and row_number are " +
+                                              "currently supported for group limits")
   }
 
   override protected def internalDoExecuteColumnar(): RDD[ColumnarBatch] = {