Define udf again

firestarman · Oct 10, 2020 · 900a91d · 900a91d
1 parent 10bd6d6
commit 900a91d
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/integration_tests/src/main/python/udf_cudf_test.py b/integration_tests/src/main/python/udf_cudf_test.py
@@ -279,6 +279,16 @@ def gpu_run(spark):
                'SpecifiedWindowFrame','UnboundedPreceding$', 'UnboundedFollowing$')
 @cudf_udf
 def test_window(enable_cudf_udf):
+    @pandas_udf("int")
+    def _sum_cpu_func(v: pd.Series) -> int:
+        return v.sum()
+
+    @pandas_udf("integer")
+    def _sum_gpu_func(v: pd.Series) -> int:
+        import cudf
+        gpu_series = cudf.Series(v)
+        return gpu_series.sum()
+
     def cpu_run(spark):
         df = _create_df(spark)
         w = Window.partitionBy('id').rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)