Fix performance regression in from_json (#10306)

* experimental fix * experimental fix * close resource * fix resource leak * signoff Signed-off-by: Andy Grove <[email protected]> --------- Signed-off-by: Andy Grove <[email protected]>
NVIDIA · Jan 29, 2024 · bc22bf8 · bc22bf8
1 parent 1fb6361
commit bc22bf8
Showing 1 changed file with 13 additions and 9 deletions.
diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuJsonToStructs.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuJsonToStructs.scala
@@ -78,19 +78,23 @@ case class GpuJsonToStructs(
               withResource(isLiteralNull.ifElse(emptyRow, nullsReplaced)) { cleaned =>
                 checkForNewline(cleaned, "\n", "line separator")
                 checkForNewline(cleaned, "\r", "carriage return")
-                // if the last entry in a column is incomplete or invalid, then cuDF
-                // will drop the row rather than replace with null if there is no newline, so we
-                // add a newline here to prevent that
-                val joined = withResource(cudf.Scalar.fromString("\n")) { lineSep =>
-                  cleaned.joinStrings(lineSep, emptyRow)
+
+                // add a newline to each JSON line
+                val withNewline = withResource(cudf.Scalar.fromString("\n")) { lineSep =>
+                  withResource(ColumnVector.fromScalar(lineSep, cleaned.getRowCount.toInt)) {
+                    newLineCol =>
+                      ColumnVector.stringConcatenate(Array[ColumnView](cleaned, newLineCol))
+                  }
                 }
-                val concat = withResource(joined) { _ =>
-                  withResource(ColumnVector.fromStrings("\n")) { newline =>
-                    ColumnVector.stringConcatenate(Array[ColumnView](joined, newline))
+
+                // join all the JSON lines into one string
+                val joined = withResource(withNewline) { _ =>
+                  withResource(Scalar.fromString("")) { emptyString =>
+                    withNewline.joinStrings(emptyString, emptyRow)
                   }
                 }
 
-                (isNullOrEmptyInput, concat)
+                (isNullOrEmptyInput, joined)
               }
             }
           }