[SPARK-59417][SQL] Stream multiline top-level JSON arrays

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JSONOptions.scala

@@ -239,6 +239,9 @@ class JSONOptions(
   val useUnsafeRow: Boolean = parameters.get(USE_UNSAFE_ROW).map(_.toBoolean).getOrElse(
     SQLConf.get.getConf(SQLConf.JSON_USE_UNSAFE_ROW))
 
+  val streamMultilineTopLevelArray: Boolean =
+    SQLConf.get.getConf(SQLConf.JSON_STREAM_MULTILINE_TOP_LEVEL_ARRAY)
+
   /** Build a Jackson [[JsonFactory]] using JSON options. */
   def buildJsonFactory(): JsonFactory = {
     val streamReadConstraints = StreamReadConstraints

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala

@@ -693,6 +693,34 @@ class JacksonParser(
     case _ => err
   }
 
+  private def badRecord(error: Throwable, recordLiteral: () => UTF8String): BadRecordException =
+    error match {
+      case e: SparkUpgradeException => throw e
+      case e: CharConversionException if options.encoding.isEmpty =>
+        val msg =
+          """JSON parser cannot handle a character in its input.
+            |Specifying encoding as an input option explicitly might help to resolve the issue.
+            |""".stripMargin + e.getMessage
+        val wrappedCharException = new CharConversionException(msg)
+        wrappedCharException.initCause(e)
+        BadRecordException(recordLiteral, () => Array.empty, wrappedCharException)
+      case PartialResultException(row, cause) =>
+        BadRecordException(recordLiteral, () => Array(row), convertCauseForPartialResult(cause))
+      case PartialResultArrayException(rows, cause) =>
+        BadRecordException(recordLiteral, () => rows, cause)
+      case PartialArrayDataResultException(arrayData, cause) =>
+        BadRecordException(
+          recordLiteral,
+          () => Array(InternalRow(arrayData)),
+          convertCauseForPartialResult(cause))
+      case PartialMapDataResultException(mapData, cause) =>
+        BadRecordException(
+          recordLiteral,
+          () => Array(InternalRow(mapData)),
+          convertCauseForPartialResult(cause))
+      case e => BadRecordException(recordLiteral, () => Array.empty, e)
+    }
+
   /**
    * Parse the JSON input to the set of [[InternalRow]]s.
    *
@@ -717,43 +745,96 @@ class JacksonParser(
       }
     } catch {
       case e: SparkUpgradeException => throw e
-      case e @ (_: RuntimeException | _: JsonProcessingException | _: MalformedInputException) =>
-        // JSON parser currently doesn't support partial results for corrupted records.
-        // For such records, all fields other than the field configured by
-        // `columnNameOfCorruptRecord` are set to `null`.
-        throw BadRecordException(() => recordLiteral(record), () => Array.empty, e)
       case e: CharConversionException if options.encoding.isEmpty =>
-        val msg =
-          """JSON parser cannot handle a character in its input.
-            |Specifying encoding as an input option explicitly might help to resolve the issue.
-            |""".stripMargin + e.getMessage
-        val wrappedCharException = new CharConversionException(msg)
-        wrappedCharException.initCause(e)
-        throw BadRecordException(() => recordLiteral(record), () => Array.empty,
-          wrappedCharException)
-      case PartialResultException(row, cause) =>
-        throw BadRecordException(
-          record = () => recordLiteral(record),
-          partialResults = () => Array(row),
-          convertCauseForPartialResult(cause))
-      case PartialResultArrayException(rows, cause) =>
-        throw BadRecordException(
-          record = () => recordLiteral(record),
-          partialResults = () => rows,
-          cause)
-      // These exceptions should never be thrown outside of JacksonParser.
-      // They are used for the control flow in the parser. We add them here for completeness
-      // since they also indicate a bad record.
-      case PartialArrayDataResultException(arrayData, cause) =>
-        throw BadRecordException(
-          record = () => recordLiteral(record),
-          partialResults = () => Array(InternalRow(arrayData)),
-          convertCauseForPartialResult(cause))
-      case PartialMapDataResultException(mapData, cause) =>
-        throw BadRecordException(
-          record = () => recordLiteral(record),
-          partialResults = () => Array(InternalRow(mapData)),
-          convertCauseForPartialResult(cause))
+        throw badRecord(e, () => recordLiteral(record))
+      case e @ (_: RuntimeException | _: JsonProcessingException | _: MalformedInputException |
+          _: PartialResultException | _: PartialResultArrayException |
+          _: PartialArrayDataResultException | _: PartialMapDataResultException) =>
+        throw badRecord(e, () => recordLiteral(record))
+    }
+  }
+
+  private[sql] def parseIterator[T](
+      record: T,
+      createParser: (JsonFactory, T) => JsonParser,
+      recordLiteral: T => UTF8String): Iterator[InternalRow] = {
+    val streamArray = allowArrayAsStructs && schema.isInstanceOf[StructType] &&
+      options.singleVariantColumn.isEmpty && options.explodeEmbeddedArray.isEmpty
+    val elementConverter = if (streamArray) makeConverter(schema) else null
+    val jsonParser = createParser(factory, record)
+    new Iterator[InternalRow] {
+      private var delegate: Iterator[InternalRow] = Iterator.empty
+      private var nextRow: InternalRow = _
+      private var prepared = false
+      private var finished = false
+      private var started = false
+      private var array = false
+
+      override def hasNext: Boolean = {
+        prepare()
+        !finished
+      }
+
+      override def next(): InternalRow = {
+        prepare()
+        if (finished) throw new NoSuchElementException("next on empty iterator")
+        prepared = false
+        nextRow
+      }
+
+      private def prepare(): Unit = {
+        if (prepared || finished) return
+        try {
+          if (!started) {
+            started = true
+            val token = jsonParser.nextToken()
+            if (token == null) {
+              finish()
+            } else if (streamArray && token == START_ARRAY) {
+              array = true
+            } else {
+              val rows = rootConverter(jsonParser)
+              if (rows == null) throw QueryExecutionErrors.rootConverterReturnNullError()
+              delegate = rows.iterator
+            }
+          }
+          if (!finished && array) {
+            jsonParser.nextToken() match {
+              case END_ARRAY => finish()
+              case null =>
+                throw new JsonParseException(jsonParser, "Unexpected end of top-level array")
+              case _ =>
+                nextRow = elementConverter(jsonParser).asInstanceOf[InternalRow]
+                if (nextRow == null) throw QueryExecutionErrors.rootConverterReturnNullError()
+                prepared = true
+            }
+          } else if (!finished && delegate.hasNext) {
+            nextRow = delegate.next()
+            prepared = true
+          } else if (!finished) {
+            finish()
+          }
+        } catch {
+          case e: SparkUpgradeException => fail(e)
+          case e: CharConversionException if options.encoding.isEmpty => fail(e)
+          case e @ (_: RuntimeException | _: JsonProcessingException | _: MalformedInputException |
+              _: PartialResultException | _: PartialResultArrayException |
+              _: PartialArrayDataResultException | _: PartialMapDataResultException) => fail(e)
+        }
+      }
+
+      private def finish(): Unit = {
+        finished = true
+        jsonParser.close()
+      }
+
+      private def fail(error: Throwable): Nothing = {
+        finished = true
+        try jsonParser.close() catch {
+          case NonFatal(closeError) => error.addSuppressed(closeError)
+        }
+        throw badRecord(error, () => recordLiteral(record))
+      }
     }
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/FailureSafeParser.scala

@@ -59,30 +59,57 @@ class FailureSafeParser[IN](
     try {
       rawParser.apply(input).iterator.map(row => toResultRow(Some(row), () => null))
     } catch {
-      case e: BadRecordException => mode match {
-        case PermissiveMode =>
-          val partialResults = e.partialResults()
-          if (partialResults.nonEmpty) {
-            partialResults.iterator.map(row => toResultRow(Some(row), e.record))
-          } else {
-            Iterator(toResultRow(None, e.record))
-          }
-        case DropMalformedMode =>
-          Iterator.empty
-        case FailFastMode =>
-          e.getCause match {
-            case _: JsonArraysAsStructsException =>
-              // SPARK-42298 we recreate the exception here to make sure the error message
-              // have the record content.
-              throw QueryExecutionErrors.cannotParseJsonArraysAsStructsError(e.record().toString)
-            case StringAsDataTypeException(fieldName, fieldValue, dataType) =>
-              throw QueryExecutionErrors.cannotParseStringAsDataTypeError(e.record().toString,
-                fieldName, fieldValue, dataType)
-            case causeWrapper: LazyBadRecordCauseWrapper =>
-              throwMalformedRecordsDetectedInRecordParsingError(e, causeWrapper.cause())
-            case cause => throwMalformedRecordsDetectedInRecordParsingError(e, cause)
-          }
+      case e: BadRecordException => parseFailure(e)
+    }
+  }
+
+  def parseIterator(
+      input: IN,
+      iteratorParser: IN => Iterator[InternalRow]): Iterator[InternalRow] = {
+    var delegate = try {
+      iteratorParser.apply(input).map(row => toResultRow(Some(row), () => null))
+    } catch {
+      case e: BadRecordException => parseFailure(e)
+    }
+    new Iterator[InternalRow] {
+      private def handleFailure[T](operation: Iterator[InternalRow] => T): T = {
+        try operation(delegate) catch {
+          case e: BadRecordException =>
+            delegate = parseFailure(e)
+            operation(delegate)
+        }
       }
+
+      override def hasNext: Boolean = handleFailure(_.hasNext)
+
+      override def next(): InternalRow = handleFailure(_.next())
+    }
+  }
+
+  private def parseFailure(e: BadRecordException): Iterator[InternalRow] = {
+    mode match {
+      case PermissiveMode =>
+        val partialResults = e.partialResults()
+        if (partialResults.nonEmpty) {
+          partialResults.iterator.map(row => toResultRow(Some(row), e.record))
+        } else {
+          Iterator(toResultRow(None, e.record))
+        }
+      case DropMalformedMode =>
+        Iterator.empty
+      case FailFastMode =>
+        e.getCause match {
+          case _: JsonArraysAsStructsException =>
+            // SPARK-42298 we recreate the exception here to make sure the error message
+            // has the record content.
+            throw QueryExecutionErrors.cannotParseJsonArraysAsStructsError(e.record().toString)
+          case StringAsDataTypeException(fieldName, fieldValue, dataType) =>
+            throw QueryExecutionErrors.cannotParseStringAsDataTypeError(e.record().toString,
+              fieldName, fieldValue, dataType)
+          case causeWrapper: LazyBadRecordCauseWrapper =>
+            throwMalformedRecordsDetectedInRecordParsingError(e, causeWrapper.cause())
+          case cause => throwMalformedRecordsDetectedInRecordParsingError(e, cause)
+        }
     }
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -7034,6 +7034,16 @@ object SQLConf {
       .booleanConf
       .createWithDefault(true)
 
+  val JSON_STREAM_MULTILINE_TOP_LEVEL_ARRAY =
+    buildConf("spark.sql.json.enableStreamingTopLevelArray")
+      .internal()
+      .doc("When true, multiline JSON reads stream the elements of a top-level array one at a " +
+        "time instead of materializing the entire array before returning rows.")
+      .version("4.4.0")
+      .withBindingPolicy(ConfigBindingPolicy.SESSION)
+      .booleanConf
+      .createWithDefault(false)
+
   val JSON_USE_UNSAFE_ROW =
     buildConf("spark.sql.json.useUnsafeRow")
       .doc("When set to true, use UnsafeRow to represent struct result in the JSON parser. It " +

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonDataSource.scala

@@ -414,8 +414,15 @@ object MultiLineJsonDataSource extends JsonDataSource {
       schema,
       parser.options.columnNameOfCorruptRecord)
 
-    safeParser.parse(
-      CodecStreams.createInputStreamWithCloseResource(conf, file.toPath))
+    val input = CodecStreams.createInputStreamWithCloseResource(conf, file.toPath)
+    Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ => input.close()))
+    if (parser.options.streamMultilineTopLevelArray) {
+      safeParser.parseIterator(
+        input,
+        input => parser.parseIterator[InputStream](input, streamParser, partitionedFileString))
+    } else {
+      safeParser.parse(input)
+    }
   }
 
   override protected def readStream(
@@ -435,6 +442,14 @@ object MultiLineJsonDataSource extends JsonDataSource {
       schema,
       parser.options.columnNameOfCorruptRecord)
 
-    safeParser.parse(new ByteArrayInputStream(bytes))
+    val input = new ByteArrayInputStream(bytes)
+    if (parser.options.streamMultilineTopLevelArray) {
+      safeParser.parseIterator(
+        input,
+        input => parser.parseIterator[InputStream](
+          input, streamParser, _ => UTF8String.fromBytes(bytes)))
+    } else {
+      safeParser.parse(input)
+    }
   }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonBenchmark.scala

@@ -17,6 +17,8 @@
 package org.apache.spark.sql.execution.datasources.json
 
 import java.io.File
+import java.nio.charset.StandardCharsets
+import java.nio.file.Files
 import java.time.{Instant, LocalDate}
 
 import org.apache.spark.benchmark.Benchmark
@@ -578,8 +580,46 @@ object JsonBenchmark extends SqlBasedBenchmark {
     benchmark.run()
   }
 
+  private def topLevelArrayBenchmark(
+      rowsNum: Int,
+      payloadSize: Int,
+      numIters: Int): Unit = {
+    val payload = "x" * payloadSize
+    val document = (0 until rowsNum)
+      .map(i => s"""{"a":$i,"payload":"$payload"}""")
+      .mkString("[", ",", "]")
+    val schema = new StructType().add("a", IntegerType).add("payload", StringType)
+    val benchmark = new Benchmark(
+      s"Top-level JSON array with $payloadSize-byte payloads", rowsNum, output = output)
+
+    withTempPath { path =>
+      Files.write(path.toPath, document.getBytes(StandardCharsets.UTF_8))
+
+      Seq(false, true).foreach { enabled =>
+        benchmark.addCase(s"streaming enabled: $enabled", numIters) { _ =>
+          withSQLConf(SQLConf.JSON_STREAM_MULTILINE_TOP_LEVEL_ARRAY.key -> enabled.toString) {
+            spark.read
+              .option("multiLine", true)
+              .schema(schema)
+              .json(path.getCanonicalPath)
+              .noop()
+          }
+        }
+      }
+
+      benchmark.run()
+    }
+  }
+
   override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
     val numIters = 3
+    if (mainArgs.contains("top-level-array")) {
+      runBenchmark("Benchmark for top-level JSON array parsing") {
+        topLevelArrayBenchmark(rowsNum = 100000, payloadSize = 0, numIters = numIters)
+        topLevelArrayBenchmark(rowsNum = 1000, payloadSize = 64 * 1024, numIters = numIters)
+      }
+      return
+    }
     runBenchmark("Benchmark for performance of JSON parsing") {
       schemaInferring(5 * 1000 * 1000, numIters)
       countShortColumn(5 * 1000 * 1000, numIters)
@@ -595,6 +635,8 @@ object JsonBenchmark extends SqlBasedBenchmark {
       // TODO (SPARK-32325): Add benchmarks for filters with nested column attributes.
       filtersPushdownBenchmark(rowsNum = 100 * 1000, numIters)
       partialResultBenchmark(rowsNum = 10000, numIters)
+      topLevelArrayBenchmark(rowsNum = 100000, payloadSize = 0, numIters = numIters)
+      topLevelArrayBenchmark(rowsNum = 1000, payloadSize = 64 * 1024, numIters = numIters)
     }
   }
 }