Add writeToStream() method to UnsafeRow

Author: JoshRosen

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java

@@ -17,6 +17,9 @@
 
 package org.apache.spark.sql.catalyst.expressions;
 
+import java.io.IOException;
+import java.io.OutputStream;
+
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.util.ObjectPool;
 import org.apache.spark.unsafe.PlatformDependent;
@@ -365,6 +368,36 @@ public InternalRow copy() {
     }
   }
 
+  /**
+   * Write this UnsafeRow's underlying bytes to the given OutputStream.
+   *
+   * @param out the stream to write to.
+   * @param writeBuffer a byte array for buffering chunks of off-heap data while writing to the
+   *                    output stream. If this row is backed by an on-heap byte array, then this
+   *                    buffer will not be used and may be null.
+   */
+  public void writeToStream(OutputStream out, byte[] writeBuffer) throws IOException {
+    if (baseObject instanceof byte[]) {
+      int offsetInByteArray = (int) (PlatformDependent.BYTE_ARRAY_OFFSET - baseOffset);
+      out.write((byte[]) baseObject, offsetInByteArray, sizeInBytes);
+    } else {
+      int dataRemaining = sizeInBytes;
+      long rowReadPosition = baseOffset;
+      while (dataRemaining > 0) {
+        int toTransfer = Math.min(writeBuffer.length, dataRemaining);
+        PlatformDependent.copyMemory(
+          baseObject,
+          rowReadPosition,
+          writeBuffer,
+          PlatformDependent.BYTE_ARRAY_OFFSET,
+          toTransfer);
+        out.write(writeBuffer, 0, toTransfer);
+        rowReadPosition += toTransfer;
+        dataRemaining -= toTransfer;
+      }
+    }
+  }
+
   @Override
   public boolean anyNull() {
     return BitSetMethods.anySet(baseObject, baseOffset, bitSetWidthInBytes);

sql/core/src/main/scala/org/apache/spark/sql/execution/UnsafeRowSerializer.scala

@@ -59,21 +59,7 @@ private class UnsafeRowSerializerInstance(numFields: Int) extends SerializerInst
       val row = value.asInstanceOf[UnsafeRow]
       assert(row.getPool == null, "UnsafeRowSerializer does not support ObjectPool")
       dOut.writeInt(row.getSizeInBytes)
-      var dataRemaining: Int = row.getSizeInBytes
-      val baseObject = row.getBaseObject
-      var rowReadPosition: Long = row.getBaseOffset
-      while (dataRemaining > 0) {
-        val toTransfer: Int = Math.min(writeBuffer.length, dataRemaining)
-        PlatformDependent.copyMemory(
-          baseObject,
-          rowReadPosition,
-          writeBuffer,
-          PlatformDependent.BYTE_ARRAY_OFFSET,
-          toTransfer)
-        out.write(writeBuffer, 0, toTransfer)
-        rowReadPosition += toTransfer
-        dataRemaining -= toTransfer
-      }
+      row.writeToStream(out, writeBuffer)
       this
     }
     override def writeKey[T: ClassTag](key: T): SerializationStream = {

sql/core/src/test/scala/org/apache/spark/sql/UnsafeRowSuite.scala

@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+import java.io.ByteArrayOutputStream
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{UnsafeRow, UnsafeProjection}
+import org.apache.spark.sql.types.{IntegerType, StringType}
+import org.apache.spark.unsafe.PlatformDependent
+import org.apache.spark.unsafe.memory.MemoryAllocator
+import org.apache.spark.unsafe.types.UTF8String
+
+class UnsafeRowSuite extends SparkFunSuite {
+  test("writeToStream") {
+    val row = InternalRow.apply(UTF8String.fromString("hello"), UTF8String.fromString("world"), 123)
+    val arrayBackedUnsafeRow: UnsafeRow =
+      UnsafeProjection.create(Seq(StringType, StringType, IntegerType)).apply(row)
+    assert(arrayBackedUnsafeRow.getBaseObject.isInstanceOf[Array[Byte]])
+    val bytesFromArrayBackedRow: Array[Byte] = {
+      val baos = new ByteArrayOutputStream()
+      arrayBackedUnsafeRow.writeToStream(baos, null)
+      baos.toByteArray
+    }
+    val bytesFromOffheapRow: Array[Byte] = {
+      val offheapRowPage = MemoryAllocator.UNSAFE.allocate(arrayBackedUnsafeRow.getSizeInBytes)
+      try {
+        PlatformDependent.copyMemory(
+          arrayBackedUnsafeRow.getBaseObject,
+          arrayBackedUnsafeRow.getBaseOffset,
+          offheapRowPage.getBaseObject,
+          offheapRowPage.getBaseOffset,
+          arrayBackedUnsafeRow.getSizeInBytes
+        )
+        val offheapUnsafeRow: UnsafeRow = new UnsafeRow()
+        offheapUnsafeRow.pointTo(
+          offheapRowPage.getBaseObject,
+          offheapRowPage.getBaseOffset,
+          3, // num fields
+          arrayBackedUnsafeRow.getSizeInBytes,
+          null // object pool
+        )
+        assert(offheapUnsafeRow.getBaseObject === null)
+        val baos = new ByteArrayOutputStream()
+        val writeBuffer = new Array[Byte](1024)
+        offheapUnsafeRow.writeToStream(baos, writeBuffer)
+        baos.toByteArray
+      } finally {
+        MemoryAllocator.UNSAFE.free(offheapRowPage)
+      }
+    }
+
+    assert(bytesFromArrayBackedRow === bytesFromOffheapRow)
+  }
+}