From 3c929c0759eeefb3a3d9c01492dbe6b594b38216 Mon Sep 17 00:00:00 2001
From: Liangcai Li <firestarmanllc@gmail.com>
Date: Thu, 13 Feb 2025 15:19:26 +0800
Subject: [PATCH] Fix an assertion error in the sized hash join (#12092)

close https://github.com/NVIDIA/spark-rapids/issues/12091

This PR tries to fix the assertion error in the sized hash join by
explicitly concatenating multiple batches into a single one.

The original code expects only one batch in the input iterator for the
build side of a small join, but it is not true when split-retry happens
in the previous operator (e.g. GpuHashAggregate), and the GPU path in
this sized join does not concatenate these small batches, instead passes
them directly down to this `getSingleBuildBatch` function. Then the
assertion fails.

One more thing, in the failing case, the join has an exchange and an
aggregate as the children.


<!--

Thank you for contributing to RAPIDS Accelerator for Apache Spark!

Here are some guidelines to help the review process go smoothly.

1. Please write a description in this text box of the changes that are
being
   made.

2. Please ensure that you have written units tests for the changes
made/features
   added.

3. If you are closing an issue please use one of the automatic closing
words as
noted here:
https://help.github.com/articles/closing-issues-using-keywords/

4. If your pull request is not ready for review but you want to make use
of the
continuous integration testing facilities please label it with `[WIP]`.

5. If your pull request is ready to be reviewed without requiring
additional
   work on top of it, then remove the `[WIP]` label (if present).

6. Once all work has been done and review has taken place please do not
add
features or make changes out of the scope of those requested by the
reviewer
(doing this just add delays as already reviewed code ends up having to
be
re-reviewed/it is hard to tell what is new etc!). Further, please avoid
rebasing your branch during the review process, as this causes the
context
of any comments made by reviewers to be lost. If conflicts occur during
review then they should be resolved by merging into the branch used for
   making the pull request.

Many thanks in advance for your cooperation!

-->

---------

Signed-off-by: Firestarman <firestarmanllc@gmail.com>
---
 .../rapids/GpuShuffledSizedHashJoinExec.scala | 23 ++++++++++++++-----
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuShuffledSizedHashJoinExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuShuffledSizedHashJoinExec.scala
index 177710fea81..50a463490bc 100644
--- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuShuffledSizedHashJoinExec.scala
+++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuShuffledSizedHashJoinExec.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -886,7 +886,7 @@ object GpuShuffledAsymmetricHashJoinExec {
           exprs.buildSideNeedsNullFilter, metrics)
         JoinInfo(joinType, buildSide, buildIter, buildSize, None, streamIter, exprs)
       } else {
-        val buildBatch = getSingleBuildBatch(baseBuildIter, exprs, metrics)
+        val buildBatch = getAsSingleBuildBatch(baseBuildIter, exprs, metrics)
         val buildIter = new SingleGpuColumnarBatchIterator(buildBatch)
         val buildStats = JoinBuildSideStats.fromBatch(buildBatch, exprs.boundBuildKeys)
         if (buildStats.streamMagnificationFactor < magnificationThreshold) {
@@ -1023,15 +1023,26 @@ object GpuShuffledAsymmetricHashJoinExec {
       }
     }
 
-    private def getSingleBuildBatch(
+    private def getAsSingleBuildBatch(
         baseIter: Iterator[ColumnarBatch],
         exprs: BoundJoinExprs,
         metrics: Map[String, GpuMetric]): ColumnarBatch = {
       val iter = addNullFilterIfNecessary(baseIter, exprs.boundBuildKeys,
         exprs.buildSideNeedsNullFilter, metrics)
-      closeOnExcept(iter.next()) { batch =>
-        assert(!iter.hasNext)
-        batch
+      // Multiple small batches may exist when split-retry happens in the previous op.
+      // So need to concat them into a single one
+      val spBatches = mutable.Queue.empty[SpillableColumnarBatch]
+      closeOnExcept(spBatches) { _ =>
+        while(iter.hasNext) {
+          spBatches.enqueue(
+            SpillableColumnarBatch(iter.next(), SpillPriorities.ACTIVE_BATCHING_PRIORITY))
+        }
+      }
+      assert(spBatches.nonEmpty, "At least one batch is expected")
+      val cbTypes = spBatches.head.dataTypes
+      withRetryNoSplit(spBatches.toSeq) { _ =>
+        ConcatAndConsumeAll.buildNonEmptyBatchFromTypes(
+          spBatches.toArray.safeMap(_.getColumnarBatch()), cbTypes)
       }
     }
   }