Add ExecChecks to Databricks shims for RunningWindowFunctionExec (NVI…

…DIA#3258) Signed-off-by: Andy Grove <[email protected]>
firestarman · Aug 19, 2021 · 0f2d575 · 0f2d575
1 parent f0f31d5
commit 0f2d575
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 9 deletions.
diff --git a/.../spark301db/src/main/scala/com/nvidia/spark/rapids/shims/spark301db/Spark301dbShims.scala b/.../spark301db/src/main/scala/com/nvidia/spark/rapids/shims/spark301db/Spark301dbShims.scala
@@ -121,10 +121,8 @@ class Spark301dbShims extends Spark301Shims {
             "i.e. (UNBOUNDED PRECEDING TO CURRENT ROW)",
         ExecChecks(
           (TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.DECIMAL_64 +
-              TypeSig.STRUCT + TypeSig.ARRAY + TypeSig.MAP).nested() +
-              TypeSig.psNote(TypeEnum.MAP, "Not supported as a partition by key") +
-              TypeSig.psNote(TypeEnum.STRUCT, "Not supported as a partition by key") +
-              TypeSig.psNote(TypeEnum.ARRAY, "Not supported as a partition by key"),
+            TypeSig.STRUCT + TypeSig.ARRAY + TypeSig.MAP).nested(),
+          Map("partitionSpec" -> (TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.DECIMAL_64)),
           TypeSig.all),
         (runningWindowFunctionExec, conf, p, r) => new GpuRunningWindowExecMeta(runningWindowFunctionExec, conf, p, r)
       ),

diff --git a/.../spark311db/src/main/scala/com/nvidia/spark/rapids/shims/spark311db/Spark311dbShims.scala b/.../spark311db/src/main/scala/com/nvidia/spark/rapids/shims/spark311db/Spark311dbShims.scala
@@ -124,11 +124,9 @@ class Spark311dbShims extends Spark311Shims {
             "i.e. (UNBOUNDED PRECEDING TO CURRENT ROW)",
         ExecChecks(
           (TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.DECIMAL_64 +
-              TypeSig.STRUCT + TypeSig.ARRAY + TypeSig.MAP).nested() +
-              TypeSig.psNote(TypeEnum.MAP, "Not supported as a partition by key") +
-              TypeSig.psNote(TypeEnum.STRUCT, "Not supported as a partition by key") +
-              TypeSig.psNote(TypeEnum.ARRAY, "Not supported as a partition by key"),
-            TypeSig.all),
+            TypeSig.STRUCT + TypeSig.ARRAY + TypeSig.MAP).nested(),
+          Map("partitionSpec" -> (TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.DECIMAL_64)),
+          TypeSig.all),
           (runningWindowFunctionExec, conf, p, r) => new GpuRunningWindowExecMeta(runningWindowFunctionExec, conf, p, r)
       ),
       GpuOverrides.exec[FileSourceScanExec](