From 31249fb59b1883110e61c4229c1333b166b2ebe4 Mon Sep 17 00:00:00 2001 From: "Hongbin Ma (Mahone)" Date: Fri, 20 Sep 2024 15:59:28 +0800 Subject: [PATCH] address comments Signed-off-by: Hongbin Ma (Mahone) --- .../scala/com/nvidia/spark/rapids/GpuExpandExec.scala | 2 ++ .../scala/com/nvidia/spark/rapids/GpuExpressions.scala | 2 +- .../nvidia/spark/rapids/basicPhysicalOperators.scala | 10 ++++++++++ 3 files changed, 13 insertions(+), 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 index afdc70a4aa7..942541d7320 100644 --- 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 index 8e7e562ed24..2fa33a597ca 100644 --- 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 index f23390849c3..891e837d7e1 100644 --- 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 {