minimize patch

Signed-off-by: Gera Shegalov <gera@apache.org>
NVIDIA · May 24, 2022 · c88cd5e · c88cd5e
1 parent dec78f4
commit c88cd5e
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/cudf_utils/HostConcatResultUtil.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/cudf_utils/HostConcatResultUtil.scala
@@ -16,7 +16,8 @@
 
 package com.nvidia.spark.rapids.cudf_utils
 
-import ai.rapids.cudf
+import ai.rapids.cudf.{HostMemoryBuffer, JCudfSerialization}
+import ai.rapids.cudf.JCudfSerialization.HostConcatResult
 import com.nvidia.spark.rapids.{Arm,  GpuColumnVectorFromBuffer}
 
 import org.apache.spark.sql.types.DataType
@@ -26,10 +27,10 @@ object HostConcatResultUtil extends Arm {
   /**
    * Create a rows-only `HostConcatResult`.
    */
-  def rowsOnlyHostConcatResult(numRows: Int): cudf.JCudfSerialization.HostConcatResult = {
-    new cudf.JCudfSerialization.HostConcatResult(
-      new cudf.JCudfSerialization.SerializedTableHeader(numRows),
-      cudf.HostMemoryBuffer.allocate(0, false))
+  def rowsOnlyHostConcatResult(numRows: Int): HostConcatResult = {
+    new HostConcatResult(
+      new JCudfSerialization.SerializedTableHeader(numRows),
+      HostMemoryBuffer.allocate(0, false))
   }
 
   /**
@@ -40,7 +41,7 @@ object HostConcatResultUtil extends Arm {
    *       callers are responsible for closing the resulting `ColumnarBatch`
    */
   def getColumnarBatch(
-      hostConcatResult: cudf.JCudfSerialization.HostConcatResult,
+      hostConcatResult: HostConcatResult,
       sparkSchema: Array[DataType]): ColumnarBatch = {
     if (hostConcatResult.getTableHeader.getNumColumns == 0) {
       // We expect the caller to have acquired the GPU unconditionally before calling