Speed up copying decimal column from parquet buffer to GPU buffer (#4872

) Signed-off-by: sperlingxx <[email protected]> Closes #4784 Adds specialized support for the columnar copy of WritableColumnVector on decimal type. The new implementation copies the unscaled values directly to avoid the round trip of Decimal encoding/decoding
NVIDIA · Mar 7, 2022 · f22795d · f22795d
1 parent 03611a5
commit f22795d
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 6 deletions.
diff --git a/sql-plugin/src/main/java/com/nvidia/spark/rapids/ColumnarCopyHelper.java b/sql-plugin/src/main/java/com/nvidia/spark/rapids/ColumnarCopyHelper.java
@@ -18,6 +18,7 @@
 
 import ai.rapids.cudf.HostColumnVector.ColumnBuilder;
 
+import org.apache.spark.sql.execution.vectorized.WritableColumnVector;
 import org.apache.spark.sql.vectorized.ColumnVector;
 
 /**
@@ -162,7 +163,30 @@ public static void stringCopy(ColumnVector cv, ColumnBuilder b, int rows) {
     }
   }
 
-  // TODO: https://github.com/NVIDIA/spark-rapids/issues/4784
+  public static void decimal32Copy(WritableColumnVector cv, ColumnBuilder b, int rows) {
+    intCopy(cv, b, rows);
+  }
+
+  public static void decimal64Copy(WritableColumnVector cv, ColumnBuilder b, int rows) {
+    longCopy(cv, b, rows);
+  }
+
+  public static void decimal128Copy(WritableColumnVector cv, ColumnBuilder b, int rows) {
+    if (!cv.hasNull()) {
+      for (int i = 0; i < rows; i++) {
+        b.appendDecimal128(cv.getBinary(i));
+      }
+      return;
+    }
+    for (int i = 0; i < rows; i++) {
+      if (cv.isNullAt(i)) {
+        b.appendNull();
+      } else {
+        b.appendDecimal128(cv.getBinary(i));
+      }
+    }
+  }
+
   public static void decimal32Copy(ColumnVector cv, ColumnBuilder b, int rows,
       int precision, int scale) {
     if (!cv.hasNull()) {

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarToGpu.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarToGpu.scala
@@ -32,6 +32,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.vectorized.WritableColumnVector
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
 import org.apache.spark.sql.vectorized.rapids.AccessibleArrowColumnVector
@@ -128,12 +129,25 @@ object HostColumnarToGpu extends Logging {
         ColumnarCopyHelper.doubleCopy(cv, b, rows)
       case StringType =>
         ColumnarCopyHelper.stringCopy(cv, b, rows)
-      case dt: DecimalType if DecimalType.is32BitDecimalType(dt) =>
-        ColumnarCopyHelper.decimal32Copy(cv, b, rows, dt.precision, dt.scale)
-      case dt: DecimalType if DecimalType.is64BitDecimalType(dt) =>
-        ColumnarCopyHelper.decimal64Copy(cv, b, rows, dt.precision, dt.scale)
       case dt: DecimalType =>
-        ColumnarCopyHelper.decimal128Copy(cv, b, rows, dt.precision, dt.scale)
+        cv match {
+          case wcv: WritableColumnVector =>
+            if (DecimalType.is32BitDecimalType(dt)) {
+              ColumnarCopyHelper.decimal32Copy(wcv, b, rows)
+            } else if (DecimalType.is64BitDecimalType(dt)) {
+              ColumnarCopyHelper.decimal64Copy(wcv, b, rows)
+            } else {
+              ColumnarCopyHelper.decimal128Copy(wcv, b, rows)
+            }
+          case _ =>
+            if (DecimalType.is32BitDecimalType(dt)) {
+              ColumnarCopyHelper.decimal32Copy(cv, b, rows, dt.precision, dt.scale)
+            } else if (DecimalType.is64BitDecimalType(dt)) {
+              ColumnarCopyHelper.decimal64Copy(cv, b, rows, dt.precision, dt.scale)
+            } else {
+              ColumnarCopyHelper.decimal128Copy(cv, b, rows, dt.precision, dt.scale)
+            }
+        }
       case t =>
         throw new UnsupportedOperationException(
           s"Converting to GPU for $t is not currently supported")