Add sequence support [databricks]

integration_tests/src/main/python/collection_ops_test.py

@@ -115,3 +115,91 @@ def test_sort_array_lit(data_gen, is_ascending):
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark: unary_op_df(spark, data_gen, length=10).select(
             f.sort_array(f.lit(array_lit), is_ascending)))
+
+# We must restrict the length of sequence, since we may suffer the exception
+# "Too long sequence: 2147483745. Should be <= 2147483632" or OOM.
+sequence_integral_gens = [
+    ByteGen(nullable=False, min_val=-20, max_val=20, special_cases=[]),
+    ShortGen(nullable=False, min_val=-20, max_val=20, special_cases=[]),
+    IntegerGen(nullable=False, min_val=-20, max_val=20, special_cases=[]),
+    LongGen(nullable=False, min_val=-20, max_val=20, special_cases=[])
+]
+
+@pytest.mark.parametrize('data_gen', sequence_integral_gens, ids=idfn)
+def test_sequence_without_step(data_gen):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark :
+        three_col_df(spark, data_gen, data_gen, data_gen)
+            .selectExpr("sequence(a, b)",
+                        "sequence(a, 0)",
+                        "sequence(0, b)"))
+
+# This function is to generate the correct sequence data according to below limitations.
+#      (step > num.zero && start <= stop)
+#        || (step < num.zero && start >= stop)
+#        || (step == num.zero && start == stop)
+def get_sequence_data(data_gen, length=2048):
+    rand = random.Random(0)
+    data_gen.start(rand)
+    list = []
+    for index in range(length):
+        start = data_gen.gen()
+        stop = data_gen.gen()
+        step = data_gen.gen()
+        # decide the direction of step
+        if start < stop:
+            step = abs(step) + 1
+        elif start == stop:
+            step = 0
+        else:
+            step = -(abs(step) + 1)
+        list.append(tuple([start, stop, step]))
+    # add special case
+    list.append(tuple([2, 2, 0]))
+    return list
+
+def get_sequence_df(spark, data, data_type):
+    return spark.createDataFrame(
+        SparkContext.getOrCreate().parallelize(data),
+        StructType([StructField('a', data_type), StructField('b', data_type), StructField('c', data_type)]))
+
+# test below case
+# (2, -1, -1)
+# (2, 5, 2)
+# (2, 2, 0)
+@pytest.mark.parametrize('data_gen', sequence_integral_gens, ids=idfn)
+def test_sequence_with_step_case1(data_gen):
+    data = get_sequence_data(data_gen)
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark :
+        get_sequence_df(spark, data, data_gen.data_type)
+            .selectExpr("sequence(a, b, c)"))
+
+sequence_three_cols_integral_gens = [
+    (ByteGen(nullable=False, min_val=-10, max_val=10, special_cases=[]),
+     ByteGen(nullable=False, min_val=30, max_val=50, special_cases=[]),
+     ByteGen(nullable=False, min_val=1, max_val=10, special_cases=[])),
+    (ShortGen(nullable=False, min_val=-10, max_val=10, special_cases=[]),
+     ShortGen(nullable=False, min_val=30, max_val=50, special_cases=[]),
+     ShortGen(nullable=False, min_val=1, max_val=10, special_cases=[])),
+    (IntegerGen(nullable=False, min_val=-10, max_val=10, special_cases=[]),
+     IntegerGen(nullable=False, min_val=30, max_val=50, special_cases=[]),
+     IntegerGen(nullable=False, min_val=1, max_val=10, special_cases=[])),
+    (LongGen(nullable=False, min_val=-10, max_val=10, special_cases=[-10, 10]),
+     LongGen(nullable=False, min_val=30, max_val=50, special_cases=[30, 50]),
+     LongGen(nullable=False, min_val=1, max_val=10, special_cases=[1, 10])),
+]
+
+# Test the scalar case for the data start < stop and step > 0
+@pytest.mark.parametrize('start_gen,stop_gen,step_gen', sequence_three_cols_integral_gens, ids=idfn)
+def test_sequence_with_step_case2(start_gen, stop_gen, step_gen):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark :
+        three_col_df(spark, start_gen, stop_gen, step_gen)
+            .selectExpr("sequence(a, b, c)",
+                        "sequence(a, b, 2)",
+                        "sequence(a, 20, c)",
+                        "sequence(a, 20, 2)",
+                        "sequence(0, b, c)",
+                        "sequence(0, 4, c)",
+                        "sequence(0, b, 3)"),)

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -3293,6 +3293,18 @@ object GpuOverrides extends Logging {
       (a, conf, p, r) => new ExprMeta[CreateMap](a, conf, p, r) {
         override def convertToGpu(): GpuExpression = GpuCreateMap(childExprs.map(_.convertToGpu()))
       }
+    ),
+    expr[Sequence](
+      desc = "Sequence",
+      ExprChecks.projectOnly(
+        TypeSig.ARRAY.nested(TypeSig.integral), TypeSig.ARRAY.nested(TypeSig.integral +
+          TypeSig.TIMESTAMP + TypeSig.DATE),
+        Seq(ParamCheck("start", TypeSig.integral, TypeSig.integral + TypeSig.TIMESTAMP +
+          TypeSig.DATE),
+          ParamCheck("stop", TypeSig.integral, TypeSig.integral + TypeSig.TIMESTAMP +
+            TypeSig.DATE)),
+        Some(RepeatingParamCheck("step", TypeSig.integral, TypeSig.integral + TypeSig.CALENDAR))),
+      (a, conf, p, r) => new GpuSequenceMeta(a, conf, p, r)
     )
   ).map(r => (r.getClassFor.asSubclass(classOf[Expression]), r)).toMap