address comments

Signed-off-by: Hongbin Ma (Mahone) <[email protected]>
NVIDIA · Sep 20, 2024 · 31249fb · 31249fb
1 parent 2984563
commit 31249fb
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpandExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpandExec.scala
@@ -66,6 +66,8 @@ class GpuExpandExecMeta(
  *                    output the same schema specified bye the parameter `output`
  * @param output      Attribute references to Output
  * @param child       Child operator
+ * @param preprojectEnabled Whether to enable pre-project before expanding
+ * @param coalesceAfter Whether to coalesce the output batches
  */
 case class GpuExpandExec(
     projections: Seq[Seq[Expression]],

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala
@@ -89,7 +89,7 @@ object GpuExpressionsUtils {
       case c: GpuColumnVector => c.incRefCount()
       case s: GpuScalar =>
         if (!s.isValid) {
-          val key = NullVecKey.apply(s.dataType, numRows)
+          val key = NullVecKey(s.dataType, numRows)
           if (!cachedNullVectors.get.containsKey(key)) {
             cachedNullVectors.get.put(key,
               GpuColumnVector.from(s, numRows, s.dataType))

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala
@@ -108,7 +108,17 @@ object GpuProjectExec {
       GpuColumnVector.incRefCounts(cb)
     } else {
       try {
+        // In some cases like Expand, we have a lot Expressions generating null vectors.
+        // We can cache the null vectors to avoid creating them every time.
+        // Since we're attempting to reuse the whole null vector, it is important to aware that
+        // datatype and vector length should be the same.
+        // Within project(cb: ColumnarBatch, boundExprs: Seq[Expression]), all output vectors share
+        // the same vector length, which facilitates the reuse of null vectors.
+        // When leaving the scope of project(cb: ColumnarBatch, boundExprs: Seq[Expression]),
+        // the cached null vectors will be cleared because the next ColumnBatch may have
+        // different vector length, thus not able to reuse cached vectors.
         GpuExpressionsUtils.cachedNullVectors.get.clear()
+
         val newColumns = boundExprs.safeMap(_.columnarEval(cb)).toArray[ColumnVector]
         new ColumnarBatch(newColumns, cb.numRows())
       } finally {