Support array_repeat

integration_tests/src/main/python/array_test.py

@@ -284,6 +284,33 @@ def test_array_max(data_gen):
                 'array_max(a)'),
             conf=no_nans_conf)
 
+
+@pytest.mark.parametrize('data_gen', orderable_gens + nested_gens_sample, ids=idfn)
+def test_array_repeat_with_count_column(data_gen):
+    cnt_gen = IntegerGen(min_val=-5, max_val=5, special_cases=[])
+    cnt_not_null_gen = IntegerGen(min_val=-5, max_val=5, special_cases=[], nullable=False)
+    gen = StructGen(
+        [('elem', data_gen), ('cnt', cnt_gen), ('cnt_nn', cnt_not_null_gen)], nullable=False)
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: gen_df(spark, gen).selectExpr(
+            'array_repeat(elem, cnt)',
+            'array_repeat(elem, cnt_nn)',
+            'array_repeat("abc", cnt)'))
+
+
+@pytest.mark.parametrize('data_gen', orderable_gens + nested_gens_sample, ids=idfn)
+def test_array_repeat_with_count_scalar(data_gen):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: unary_op_df(spark, data_gen).selectExpr(
+            'array_repeat(a, 3)',
+            'array_repeat(a, 1)',
+            'array_repeat(a, 0)',
+            'array_repeat(a, -2)',
+            'array_repeat("abc", 2)',
+            'array_repeat("abc", 0)',
+            'array_repeat("abc", -1)'))
+
+
 # We add in several types of processing for foldable functions because the output
 # can be different types.
 @pytest.mark.parametrize('query', [

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

@@ -2768,6 +2768,20 @@ object GpuOverrides extends Logging {
         override def convertToGpu(child: Expression): GpuExpression =
           GpuArrayMax(child)
       }),
+    expr[ArrayRepeat](
+      "Returns the array containing the given input value (left) count (right) times",
+      ExprChecks.binaryProject(
+        TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL
+          + TypeSig.ARRAY + TypeSig.STRUCT + TypeSig.MAP),
+        TypeSig.ARRAY.nested(TypeSig.all),
+        ("left", (TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL
+          + TypeSig.ARRAY + TypeSig.STRUCT + TypeSig.MAP).nested(), TypeSig.all),
+        ("right", TypeSig.integral, TypeSig.integral)),
+      (in, conf, p, r) => new BinaryExprMeta[ArrayRepeat](in, conf, p, r) {
+        override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression =
+          GpuArrayRepeat(lhs, rhs)
+      }
+    ),
     expr[CreateNamedStruct](
       "Creates a struct with the given field names and values",
       CreateNamedStructCheck,

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/collectionOperations.scala

@@ -508,6 +508,91 @@ case class GpuArrayMax(child: Expression) extends GpuUnaryExpression with Implic
     }
 }
 
+case class GpuArrayRepeat(left: Expression, right: Expression) extends GpuBinaryExpression {
+
+  override def dataType: DataType = ArrayType(left.dataType, left.nullable)
+
+  override def doColumnar(lhs: GpuColumnVector, rhs: GpuColumnVector): ColumnVector = {
+    // The primary issue of array_repeat is to workaround the null and negative count.
+    // Spark returns a null (list) when encountering a null count, while cudf::repeat simply
+    // throws an exception.
+    // Spark returns a empty list when encountering a negative count, while cudf::repeat simply
+    // throws an exception.
+
+    // Step 1. replace invalid counts
+    //  null -> 1
+    //  negative values -> 0
+    val refinedCount = withResource(GpuScalar.from(1, DataTypes.IntegerType)) { one =>
+      withResource(rhs.getBase.replaceNulls(one)) { notNull =>
+        withResource(GpuScalar.from(0, DataTypes.IntegerType)) { zero =>
+          withResource(notNull.lessThan(zero)) { lessThanZero =>
+            lessThanZero.ifElse(zero, notNull)
+          }
+        }
+      }
+    }
+    // Step 2. perform cuDF repeat
+    val repeated = closeOnExcept(refinedCount) { cnt =>
+      withResource(new Table(lhs.getBase)) { table =>
+        table.repeat(cnt, true).getColumn(0)
+      }
+    }
+    // Step 3. generate list offsets from refined counts
+    val offsets = withResource(refinedCount) { cnt =>
+      cnt.generateListOffsets()
+    }
+    // Step 4. make the result list column with offsets and child column
+    val list = withResource(offsets) { offsets =>
+      withResource(repeated) { repeated =>
+        repeated.makeListFromOffsets(lhs.getRowCount, offsets)
+      }
+    }
+    // Step 5. merge the validity of count column to the result
+    withResource(list) { list =>
+      list.mergeAndSetValidity(BinaryOp.BITWISE_AND, rhs.getBase)
+    }
+  }
+
+  override def doColumnar(lhs: GpuScalar, rhs: GpuColumnVector): ColumnVector = {
+    withResource(GpuColumnVector.from(lhs, rhs.getRowCount.toInt, lhs.dataType)) { left =>
+      doColumnar(left, rhs)
+    }
+  }
+
+  override def doColumnar(lhs: GpuColumnVector, rhs: GpuScalar): ColumnVector = {
+    val numRows = lhs.getRowCount.toInt
+    if (!rhs.isValid) {
+      GpuColumnVector.fromNull(numRows, dataType).getBase
+    } else {
+      val count = rhs.getValue.asInstanceOf[Int] max 0
+
+      val offsets = withResource(GpuScalar.from(count, IntegerType)) { cntScalar =>
+        withResource(GpuColumnVector.from(cntScalar, numRows, rhs.dataType)) { cnt =>
+          cnt.getBase.generateListOffsets()
+        }
+      }
+
+      withResource(offsets) { offsets =>
+        withResource(new Table(lhs.getBase)) { table =>
+          withResource(table.repeat(count).getColumn(0)) { repeated =>
+            repeated.makeListFromOffsets(lhs.getRowCount, offsets)
+          }
+        }
+      }
+    }
+  }
+
+  override def doColumnar(numRows: Int, lhs: GpuScalar, rhs: GpuScalar): ColumnVector = {
+    if (!rhs.isValid) {
+      GpuColumnVector.fromNull(numRows, dataType).getBase
+    } else {
+      withResource(GpuColumnVector.from(lhs, numRows, lhs.dataType)) { left =>
+        doColumnar(left, rhs)
+      }
+    }
+  }
+}
+
 class GpuSequenceMeta(
     expr: Sequence,
     conf: RapidsConf,