[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,10 @@ class JSONOptions(
   val useUnsafeRow: Boolean = parameters.get(USE_UNSAFE_ROW).map(_.toBoolean).getOrElse(
     SQLConf.get.getConf(SQLConf.JSON_USE_UNSAFE_ROW))
 
+  val streamMultilineTopLevelArray: Boolean =
+    parameters.get(ENABLE_STREAMING_TOP_LEVEL_ARRAY).map(_.toBoolean).getOrElse(
+      SQLConf.get.getConf(SQLConf.JSON_STREAM_MULTILINE_TOP_LEVEL_ARRAY))
+
   /** Build a Jackson [[JsonFactory]] using JSON options. */
   def buildJsonFactory(): JsonFactory = {
     val streamReadConstraints = StreamReadConstraints
@@ -338,6 +342,7 @@ object JSONOptions extends DataSourceOptions {
   val SINGLE_VARIANT_COLUMN = newOption(DataSourceOptions.SINGLE_VARIANT_COLUMN)
   val EXPLODE_EMBEDDED_ARRAY = newOption(DataSourceOptions.EXPLODE_EMBEDDED_ARRAY)
   val USE_UNSAFE_ROW = newOption("useUnsafeRow")
+  val ENABLE_STREAMING_TOP_LEVEL_ARRAY = newOption("enableStreamingTopLevelArray")
   // Options with alternative
   val ENCODING = "encoding"
   val CHARSET = "charset"

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

@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.json
 
 import java.io.{ByteArrayOutputStream, CharConversionException}
 import java.nio.charset.MalformedInputException
+import java.util.concurrent.atomic.AtomicReference
 
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
@@ -28,6 +29,7 @@ import com.fasterxml.jackson.core._
 import org.apache.hadoop.fs.PositionedReadable
 
 import org.apache.spark.SparkUpgradeException
+import org.apache.spark.TaskContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.{InternalRow, NoopFilters, StructFilters}
 import org.apache.spark.sql.catalyst.expressions._
@@ -61,6 +63,9 @@ class JacksonParser(
   // `ValueConverter`s for the root schema for all fields in the schema
   private val rootConverter = makeRootConverter(schema)
 
+  // Lazy: a schema that is not a struct never reaches either top-level-array path.
+  private lazy val arrayElementConverter = makeConverter(schema)
+
   private val factory = options.buildJsonFactory()
 
   private lazy val timestampFormatter = TimestampFormatter(
@@ -148,7 +153,6 @@ class JacksonParser(
   }
 
   private def makeStructRootConverter(st: StructType): JsonParser => Iterable[InternalRow] = {
-    val elementConverter = makeConverter(st)
     val fieldConverters = st.map(_.dataType).map(makeConverter).toArray
     val jsonFilters = if (SQLConf.get.jsonFilterPushDown) {
       new JsonFilters(filters, st)
@@ -171,7 +175,8 @@ class JacksonParser(
         // List([str_a_2,null], [null,str_b_3])
         //
       case START_ARRAY if allowArrayAsStructs =>
-        val array = convertArray(parser, elementConverter, isRoot = true, arrayAsStructs = true)
+        val array =
+          convertArray(parser, arrayElementConverter, isRoot = true, arrayAsStructs = true)
         // Here, as we support reading top level JSON arrays and take every element
         // in such an array as a row, this case is possible.
         if (array.numElements() == 0) {
@@ -693,6 +698,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 +750,122 @@ 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 jsonParser = try {
+      createParser(factory, record)
+    } catch {
+      case e: SparkUpgradeException => throw e
+      case e: CharConversionException if options.encoding.isEmpty =>
+        throw badRecord(e, () => recordLiteral(record))
+      case e @ (_: RuntimeException | _: JsonProcessingException | _: MalformedInputException |
+          _: PartialResultException | _: PartialResultArrayException |
+          _: PartialArrayDataResultException | _: PartialMapDataResultException) =>
+        throw badRecord(e, () => recordLiteral(record))
+    }
+    // An archive read builds one parser per entry inside a single task, so a closed parser left
+    // reachable from the completion listener keeps its entry's buffer alive for the rest of the
+    // task. Every close path takes the parser out of this cell first, so exactly one of them closes
+    // it and the listener stops retaining it as soon as one does.
+    val openParser = new AtomicReference(jsonParser)
+    def closeParser(): Unit = Option(openParser.getAndSet(null)).foreach(_.close())
+    // Abandoning the iterator before END_ARRAY (e.g. a LIMIT) reaches no eager close.
+    Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ => closeParser()))
+    def fail(error: Throwable): Nothing = {
+      try closeParser() catch {
+        case NonFatal(closeError) => error.addSuppressed(closeError)
+      }
+      throw badRecord(error, () => recordLiteral(record))
+    }
+    // Resuming is only safe where the parser cannot have moved past the element that failed. A
+    // scalar element leaves it on that scalar's own token, and `convertObject` consumes through
+    // END_OBJECT before raising `PartialResultException`; any other container failure can have
+    // stopped anywhere inside the element.
+    def resumableAfter(error: Throwable, elementStart: Option[JsonToken]): Boolean =
+      options.parseMode != FailFastMode && (elementStart match {
+        case Some(START_OBJECT) | Some(START_ARRAY) => error.isInstanceOf[PartialResultException]
+        case Some(_) => error.isInstanceOf[RuntimeException]
+        case None => false
+      })
+    def handleFailure[R](elementStart: Option[JsonToken])(operation: => R): R = {
+      try operation catch {
+        case e: SparkUpgradeException => fail(e)
+        case e: CharConversionException if options.encoding.isEmpty => fail(e)
+        case e @ (_: RuntimeException | _: JsonProcessingException | _: MalformedInputException |
+            _: PartialResultException | _: PartialResultArrayException |
+            _: PartialArrayDataResultException | _: PartialMapDataResultException) =>
+          if (resumableAfter(e, elementStart)) {
+            throw badRecord(e, () => recordLiteral(record)).copy(recoverable = true)
+          } else {
+            fail(e)
+          }
+      }
+    }
+
+    handleFailure(elementStart = None)(jsonParser.nextToken()) match {
+      case null =>
+        closeParser()
+        Iterator.empty
+      case START_ARRAY if streamArray =>
+        new Iterator[InternalRow] {
+          private var nextRow: InternalRow = _
+          private var prepared = false
+          private var finished = 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
+            handleFailure(elementStart = None)(jsonParser.nextToken()) match {
+              case END_ARRAY => finish()
+              // Unreachable today: jackson-core signals end of input inside an open array by
+              // throwing JsonEOFException, which handleFailure turns into a malformed record.
+              // Kept so a backend that returns null cannot fall through with no current token.
+              case null =>
+                fail(new JsonParseException(jsonParser, "Unexpected end of top-level array"))
+              case elementStart =>
+                nextRow = handleFailure(Some(elementStart)) {
+                  val row = arrayElementConverter(jsonParser).asInstanceOf[InternalRow]
+                  if (row == null) throw QueryExecutionErrors.rootConverterReturnNullError()
+                  row
+                }
+                prepared = true
+            }
+          }
+
+          private def finish(): Unit = {
+            finished = true
+            closeParser()
+          }
+        }
+      case _ =>
+        val rows = handleFailure(elementStart = None)(rootConverter(jsonParser))
+        if (rows == null) fail(QueryExecutionErrors.rootConverterReturnNullError())
+        closeParser()
+        rows.iterator
     }
   }
 }

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

@@ -76,11 +76,13 @@ case class PartialResultArrayException(
  * @param cause the actual exception about why the record is bad and can't be parsed. It's better
  *                      to use `LazyBadRecordCauseWrapper` here to delay heavy cause construction
  *                      until it's needed.
+ * @param recoverable whether the parser can resume at the next record after handling this failure
  */
 case class BadRecordException(
     @transient record: () => UTF8String,
     @transient partialResults: () => Array[InternalRow] = () => Array.empty[InternalRow],
-    cause: Throwable) extends Exception(cause) {
+    cause: Throwable,
+    recoverable: Boolean = false) extends Exception(cause) {
   override def getStackTrace(): Array[StackTraceElement] = new Array[StackTraceElement](0)
   override def fillInStackTrace(): Throwable = this
 }

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

@@ -59,30 +59,63 @@ 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 = {
+        while (true) {
+          try {
+            return operation(delegate)
+          } catch {
+            case e: BadRecordException =>
+              val source = delegate
+              val recovery = parseFailure(e)
+              delegate = if (e.recoverable) recovery ++ source else recovery
           }
+        }
+        throw new IllegalStateException("unreachable")
       }
+
+      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,25 @@ 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. This applies only " +
+        "to reads into a struct schema that take top-level arrays as structs; reads using the " +
+        "`singleVariantColumn` or `explodeEmbeddedArray` option are never streamed. Streaming " +
+        "also makes an array element, rather than the whole document, the record that a parse " +
+        "mode applies to, since rows already emitted cannot be withdrawn: PERMISSIVE fills the " +
+        "corrupt record column for the malformed element only, leaving it null on the valid rows " +
+        "of the same document, and DROPMALFORMED drops that element rather than the whole " +
+        "document. An element whose failure leaves the parser at an unknown position, such as a " +
+        "nested value of the wrong shape, still ends the document. It can be overwritten by the " +
+        "JSON option `enableStreamingTopLevelArray`.")
+      .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 " +