Apply code review suggestion

NVIDIA · Jan 11, 2024 · 123b5d8 · 123b5d8
1 parent 9e7abd1
commit 123b5d8
Showing 1 changed file with 7 additions and 8 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuGenerateExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuGenerateExec.scala
@@ -550,14 +550,13 @@ abstract class GpuExplodeBase extends GpuUnevaluableUnaryExpression with GpuGene
           posColToFix.add(offset, posColToFix.getType)
         }
         withResource(fixedUpPosCol) { _ =>
-          val newCols = new Array[ColumnVector](toFixUp.getNumberOfColumns)
-          (0 until genOffset).foreach { b =>
-            newCols(b) = toFixUp.getColumn(b)
-          }
-          newCols(genOffset) = fixedUpPosCol
-          (genOffset + 1 until toFixUp.getNumberOfColumns).foreach { a =>
-            newCols(a) = toFixUp.getColumn(a)
-          }
+          val newCols = (0 until toFixUp.getNumberOfColumns).map { i =>
+            if (i == genOffset) {
+              fixedUpPosCol
+            } else {
+              toFixUp.getColumn(i)
+            }
+          }.toArray
           new Table(newCols: _*)
         }
       }