diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/file/FileScanUtils.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/file/FileScanUtils.scala index a7f81b4869c..6ae7a4632bc 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/file/FileScanUtils.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/file/FileScanUtils.scala @@ -21,7 +21,7 @@ package org.apache.texera.amber.operator.source.scan.file import org.apache.commons.compress.archivers.ArchiveStreamFactory import org.apache.commons.compress.archivers.zip.ZipArchiveInputStream -import org.apache.commons.io.IOUtils.toByteArray + import org.apache.texera.amber.core.storage.DocumentFactory import org.apache.texera.amber.core.tuple.AttributeTypeUtils.parseField import org.apache.texera.amber.core.tuple.LargeBinary @@ -39,6 +39,34 @@ import scala.collection.mutable import scala.jdk.CollectionConverters.IteratorHasAsScala private[file] object FileScanUtils { + + private[file] def safeToByteArray( + entry: InputStream, + attributeType: FileAttributeType + ): Array[Byte] = { + try { + val out = new ByteArrayOutputStream() + val buffer = new Array[Byte](8192) + var bytesRead = entry.read(buffer) + + while (bytesRead != -1) { + out.write(buffer, 0, bytesRead) + bytesRead = entry.read(buffer) + } + out.toByteArray + } catch { + case _: OutOfMemoryError | _: IllegalArgumentException => + val largeBinaryHint = attributeType match { + case FileAttributeType.BINARY => "Please use 'large binary' attribute type instead." + case FileAttributeType.SINGLE_STRING => + "Please split the file or use a chunked reading method." + case _ => "File is too large to fit in memory." + } + throw new RuntimeException( + s"File exceeds maximum safe memory size for '${attributeType.getName}' type. $largeBinaryHint" + ) + } + } def createTuplesFromFile( fileName: String, displayFileName: String, @@ -90,7 +118,7 @@ private[file] object FileScanUtils { } fields += (attributeType match { case FileAttributeType.SINGLE_STRING => - new String(toByteArray(entry), fileEncoding.getCharset) + new String(safeToByteArray(entry, attributeType), fileEncoding.getCharset) case FileAttributeType.LARGE_BINARY => val largeBinary = new LargeBinary() val out = new LargeBinaryOutputStream(largeBinary) @@ -105,7 +133,7 @@ private[file] object FileScanUtils { out.close() } largeBinary - case _ => parseField(toByteArray(entry), attributeType.getType) + case _ => parseField(safeToByteArray(entry, attributeType), attributeType.getType) }) TupleLike(fields.toSeq: _*) } diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/file/FileScanUtilsSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/file/FileScanUtilsSpec.scala index ad2ffeb79e6..507c6d81e9c 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/file/FileScanUtilsSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/file/FileScanUtilsSpec.scala @@ -105,4 +105,20 @@ class FileScanUtilsSpec extends AnyFlatSpec with BeforeAndAfterAll { .toSeq assert(contents(tuples) == Seq("l1", "l2", "l3")) } + it should "throw RuntimeException when binary file exceeds natural memory limit" in { + val mockLargeInputStream = new java.io.InputStream { + override def read(): Int = { + throw new OutOfMemoryError("Requested array size exceeds VM limit") + } + override def read(b: Array[Byte], off: Int, len: Int): Int = { + throw new OutOfMemoryError("Requested array size exceeds VM limit") + } + } + + val exception = intercept[RuntimeException] { + FileScanUtils.safeToByteArray(mockLargeInputStream, FileAttributeType.BINARY) + } + assert(exception.getMessage.contains("exceeds maximum safe memory size")) + assert(exception.getMessage.contains("Please use 'large binary'")) + } }