Fix batch data loading in the ByteReader

JD557 · Mar 1, 2024 · ade1823 · ade1823
1 parent 414f21e
commit ade1823
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/core/shared/src/main/scala/eu/joaocosta/minart/internal/ByteReader.scala b/core/shared/src/main/scala/eu/joaocosta/minart/internal/ByteReader.scala
@@ -30,14 +30,22 @@ private[minart] object ByteReader {
   /** Read N Bytes */
   def readBytes(n: Int): ParseState[Nothing, Array[Int]] = State { bytes =>
     val byteArr = Array.ofDim[Byte](n)
-    bytes.read(byteArr)
+    var read    = bytes.read(byteArr)
+    while (read >= 0 && read < n) {
+      byteArr(read) = bytes.read().toByte
+      read += 1
+    }
     bytes -> byteArr.map(b => java.lang.Byte.toUnsignedInt(b))
   }
 
   /** Read N Bytes */
   def readRawBytes(n: Int): ParseState[Nothing, Array[Byte]] = State { bytes =>
     val byteArr = Array.ofDim[Byte](n)
-    bytes.read(byteArr)
+    var read    = bytes.read(byteArr)
+    while (read >= 0 && read < n) {
+      byteArr(read) = bytes.read().toByte
+      read += 1
+    }
     bytes -> byteArr
   }