From 4d7cbfdf3e3819946a560d312b1e1cb9a0e659da Mon Sep 17 00:00:00 2001
From: Yan Feng <fengyan_@mail.ustc.edu.cn>
Date: Thu, 12 Dec 2024 16:43:25 +0800
Subject: [PATCH] Retrieve supportedDepth from JNI

Signed-off-by: Yan Feng <fengyan_@mail.ustc.edu.cn>
---
 docs/supported_ops.md                         |  4 ++--
 .../nvidia/spark/rapids/GpuOverrides.scala    | 20 ++++++++++---------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/docs/supported_ops.md b/docs/supported_ops.md
index 30dbc606b20..be52afd9799 100644
--- a/docs/supported_ops.md
+++ b/docs/supported_ops.md
@@ -8480,9 +8480,9 @@ are limited.
 <td>S</td>
 <td><b>NS</b></td>
 <td><b>NS</b></td>
-<td><em>PS<br/>Nested levels exceeding 8 layers are not supported;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types DECIMAL, BINARY, CALENDAR, MAP, UDT, DAYTIME, YEARMONTH</em></td>
+<td><em>PS<br/>The nesting depth has a certain limit;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types DECIMAL, BINARY, CALENDAR, MAP, UDT, DAYTIME, YEARMONTH</em></td>
 <td><b>NS</b></td>
-<td><em>PS<br/>Nested levels exceeding 8 layers are not supported;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types DECIMAL, BINARY, CALENDAR, MAP, UDT, DAYTIME, YEARMONTH</em></td>
+<td><em>PS<br/>The nesting depth has a certain limit;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types DECIMAL, BINARY, CALENDAR, MAP, UDT, DAYTIME, YEARMONTH</em></td>
 <td><b>NS</b></td>
 <td><b>NS</b></td>
 <td><b>NS</b></td>
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 a687a90d4c5..bd9aad97772 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
@@ -25,6 +25,7 @@ import scala.util.control.NonFatal
 import ai.rapids.cudf.DType
 import com.nvidia.spark.rapids.RapidsConf.{SUPPRESS_PLANNING_FAILURE, TEST_CONF}
 import com.nvidia.spark.rapids.jni.GpuTimeZoneDB
+import com.nvidia.spark.rapids.jni.Hash
 import com.nvidia.spark.rapids.lore.GpuLore
 import com.nvidia.spark.rapids.shims._
 import com.nvidia.spark.rapids.window.{GpuDenseRank, GpuLag, GpuLead, GpuPercentRank, GpuRank, GpuRowNumber, GpuSpecialFrameBoundary, GpuWindowExecMeta, GpuWindowSpecDefinitionMeta}
@@ -3321,23 +3322,24 @@ object GpuOverrides extends Logging {
       ExprChecks.projectOnly(TypeSig.INT, TypeSig.INT,
         repeatingParamCheck = Some(RepeatingParamCheck("input",
           (TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.STRUCT + TypeSig.ARRAY).nested() +
-              TypeSig.psNote(TypeEnum.ARRAY, "Nested levels exceeding 8 layers are not supported") +
-              TypeSig.psNote(TypeEnum.STRUCT, "Nested levels exceeding 8 layers are not supported"),
+              TypeSig.psNote(TypeEnum.ARRAY, "The nesting depth has a certain limit") +
+              TypeSig.psNote(TypeEnum.STRUCT, "The nesting depth has a certain limit"),
           TypeSig.all))),
       (a, conf, p, r) => new ExprMeta[HiveHash](a, conf, p, r) {
         override def tagExprForGpu(): Unit = {
-          def getMaxNestedDepth(inputType: DataType): Int = {
+          def getMaxStackDepth(inputType: DataType): Int = {
             inputType match {
-              case at: ArrayType => 1 + getMaxNestedDepth(at.elementType)
+              case at: ArrayType => 1 + getMaxStackDepth(at.elementType)
               case st: StructType =>
-                1 + st.map(f => getMaxNestedDepth(f.dataType)).max
+                1 + st.map(f => getMaxStackDepth(f.dataType)).max
               case _ => 0 // primitive types
             }
           }
-          val maxDepth = a.children.map(c => getMaxNestedDepth(c.dataType)).max
-          if (maxDepth > 8) {
-            willNotWorkOnGpu(s"GPU HiveHash supports 8 levels at most for " +
-              s"nested types, but got $maxDepth")
+          val maxDepth = a.children.map(c => getMaxStackDepth(c.dataType)).max
+          val supportedDepth = Hash.MAX_STACK_DEPTH
+          if (maxDepth > supportedDepth) {
+            willNotWorkOnGpu(s"the data type requires a stack size of $maxDepth, " +
+              s"which exceeds the GPU limit of $supportedDepth")
           }
         }