Handle AggregateExpression having resultIds parameter instead of a (#189

) single resultId Co-authored-by: Thomas Graves <[email protected]>
NVIDIA · Jun 17, 2020 · 10a3ecb · 10a3ecb
1 parent c50fcbd
commit 10a3ecb
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/sql-plugin/src/main/scala/ai/rapids/spark/GpuOverrides.scala b/sql-plugin/src/main/scala/ai/rapids/spark/GpuOverrides.scala
@@ -1136,9 +1136,20 @@ object GpuOverrides {
         } else {
           childrenExprMeta
         }
-        override def convertToGpu(): GpuExpression =
-          GpuAggregateExpression(childExprs.head.convertToGpu().asInstanceOf[GpuAggregateFunction],
-            a.mode, a.isDistinct, filter.map(_.convertToGpu()) ,a.resultId)
+       override def convertToGpu(): GpuExpression = {
+         // handle the case AggregateExpression has the resultIds parameter where its
+         // Seq[ExprIds] instead of single ExprId.
+         val resultId = try {
+            val resultMethod = a.getClass.getMethod("resultId")
+            resultMethod.invoke(a).asInstanceOf[ExprId]
+          } catch {
+            case e: Exception =>
+              val resultMethod = a.getClass.getMethod("resultIds")
+              resultMethod.invoke(a).asInstanceOf[Seq[ExprId]](0)
+          }
+          GpuAggregateExpression(childExprs(0).convertToGpu().asInstanceOf[GpuAggregateFunction],
+            a.mode, a.isDistinct, filter.map(_.convertToGpu()) ,resultId)
+       }
       }),
     expr[SortOrder](
       "sort order",