Filter out the nulls after slicing the batches. (#2447)

* Filter out the nulls after slicing the batches. Signed-off-by: Firestarman <firestarmanllc@gmail.com>
NVIDIA · May 20, 2021 · 2b781d8 · 2b781d8
1 parent 6874645
commit 2b781d8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuRangePartitioner.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuRangePartitioner.scala
@@ -210,7 +210,7 @@ case class GpuRangePartitioner(
         // slice will produce.
         val sliced = sliceInternalGpuOrCpu(cb.numRows, Array(0),
           GpuColumnVector.extractColumns(cb))
-        sliced.zipWithIndex
+        sliced.zipWithIndex.filter(_._1 != null)
       }
     }
   }

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuSinglePartitioning.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuSinglePartitioning.scala
@@ -44,7 +44,7 @@ case object GpuSinglePartitioning extends GpuExpression with GpuPartitioning {
           batch.numRows,
           Array(0),
           GpuColumnVector.extractColumns(batch))
-        sliced.zipWithIndex
+        sliced.zipWithIndex.filter(_._1 != null)
       }
     }
   }