diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala index 20f52997ab2..341b11c4619 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala @@ -63,6 +63,8 @@ import javax.ws.rs.core.UriBuilder import java.util.Locale +import scala.collection.JavaConverters._ + class VeloxSparkPlanExecApi extends SparkPlanExecApi { /** Transform GetArrayItem to Substrait. */ @@ -670,26 +672,32 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi { dataSize: SQLMetric): BuildSideRelation = { val useOffheapBroadcastBuildRelation = VeloxConfig.get.enableBroadcastBuildRelationInOffheap - val serialized: Array[ColumnarBatchSerializeResult] = child + val serialized: Seq[ColumnarBatchSerializeResult] = child .executeColumnar() .mapPartitions(itr => Iterator(BroadcastUtils.serializeStream(itr))) - .filter(_.getNumRows != 0) + .filter(_.numRows != 0) .collect - val rawSize = serialized.flatMap(_.getSerialized.map(_.length.toLong)).sum + val rawSize = serialized.map(_.sizeInBytes()).sum if (rawSize >= GlutenConfig.get.maxBroadcastTableSize) { throw new SparkException( "Cannot broadcast the table that is larger than " + s"${SparkMemoryUtil.bytesToString(GlutenConfig.get.maxBroadcastTableSize)}: " + s"${SparkMemoryUtil.bytesToString(rawSize)}") } - numOutputRows += serialized.map(_.getNumRows).sum + numOutputRows += serialized.map(_.numRows).sum dataSize += rawSize if (useOffheapBroadcastBuildRelation) { TaskResources.runUnsafe { - UnsafeColumnarBuildSideRelation(child.output, serialized.flatMap(_.getSerialized), mode) + UnsafeColumnarBuildSideRelation( + child.output, + serialized.flatMap(_.offHeapData().asScala), + mode) } } else { - ColumnarBuildSideRelation(child.output, serialized.flatMap(_.getSerialized), mode) + ColumnarBuildSideRelation( + child.output, + serialized.flatMap(_.onHeapData().asScala).toArray, + mode) } } diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala index 5fc9cc4464e..ad066d47f9e 100644 --- a/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala +++ b/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala @@ -34,6 +34,7 @@ import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.task.TaskResources +import scala.collection.JavaConverters._ import scala.collection.mutable.ArrayBuffer // Utility methods to convert Vanilla broadcast relations from/to Velox broadcast relations. @@ -93,58 +94,44 @@ object BroadcastUtils { schema: StructType, from: Broadcast[F], fn: Iterator[InternalRow] => Iterator[ColumnarBatch]): Broadcast[T] = { - val useOffheapBuildRelation = VeloxConfig.get.enableBroadcastBuildRelationInOffheap + + def batchIterationToRelation(batchItr: () => Iterator[ColumnarBatch]): BuildSideRelation = { + TaskResources.runUnsafe { + serializeStream(batchItr()) match { + case ColumnarBatchSerializeResult.EMPTY => + ColumnarBuildSideRelation( + SparkShimLoader.getSparkShims.attributesFromStruct(schema), + Array[Array[Byte]](), + mode) + case result: ColumnarBatchSerializeResult => + if (result.isOffHeap) { + UnsafeColumnarBuildSideRelation( + SparkShimLoader.getSparkShims.attributesFromStruct(schema), + result.offHeapData().asScala.toSeq, + mode) + } else { + ColumnarBuildSideRelation( + SparkShimLoader.getSparkShims.attributesFromStruct(schema), + result.onHeapData().asScala.toArray, + mode) + } + } + } + } + mode match { case HashedRelationBroadcastMode(_, _) => // HashedRelation to ColumnarBuildSideRelation. val fromBroadcast = from.asInstanceOf[Broadcast[HashedRelation]] val fromRelation = fromBroadcast.value.asReadOnlyCopy() - val toRelation = TaskResources.runUnsafe { - val batchItr: Iterator[ColumnarBatch] = fn(reconstructRows(fromRelation)) - val serialized: Array[Array[Byte]] = serializeStream(batchItr) match { - case ColumnarBatchSerializeResult.EMPTY => - Array() - case result: ColumnarBatchSerializeResult => - result.getSerialized - } - if (useOffheapBuildRelation) { - UnsafeColumnarBuildSideRelation( - SparkShimLoader.getSparkShims.attributesFromStruct(schema), - serialized, - mode) - } else { - ColumnarBuildSideRelation( - SparkShimLoader.getSparkShims.attributesFromStruct(schema), - serialized, - mode) - } - } + val toRelation = batchIterationToRelation(() => fn(reconstructRows(fromRelation))) // Rebroadcast Velox relation. context.broadcast(toRelation).asInstanceOf[Broadcast[T]] case IdentityBroadcastMode => // Array[InternalRow] to ColumnarBuildSideRelation. val fromBroadcast = from.asInstanceOf[Broadcast[Array[InternalRow]]] val fromRelation = fromBroadcast.value - val toRelation = TaskResources.runUnsafe { - val batchItr: Iterator[ColumnarBatch] = fn(fromRelation.iterator) - val serialized: Array[Array[Byte]] = serializeStream(batchItr) match { - case ColumnarBatchSerializeResult.EMPTY => - Array() - case result: ColumnarBatchSerializeResult => - result.getSerialized - } - if (useOffheapBuildRelation) { - UnsafeColumnarBuildSideRelation( - SparkShimLoader.getSparkShims.attributesFromStruct(schema), - serialized, - mode) - } else { - ColumnarBuildSideRelation( - SparkShimLoader.getSparkShims.attributesFromStruct(schema), - serialized, - mode) - } - } + val toRelation = batchIterationToRelation(() => fn(fromRelation.iterator)) // Rebroadcast Velox relation. context.broadcast(toRelation).asInstanceOf[Broadcast[T]] case _ => throw new IllegalStateException("Unexpected broadcast mode: " + mode) @@ -175,25 +162,25 @@ object BroadcastUtils { val handle = ColumnarBatches.getNativeHandle(BackendsApiManager.getBackendName, b) numRows += b.numRows() try { - val unsafeBuffer = ColumnarBatchSerializerJniWrapper + ColumnarBatchSerializerJniWrapper .create( Runtimes .contextInstance( BackendsApiManager.getBackendName, "BroadcastUtils#serializeStream")) .serialize(handle) - try { - unsafeBuffer.toByteArray - } finally { - unsafeBuffer.close() - } } finally { ColumnarBatches.release(b) } }) .toArray if (values.nonEmpty) { - new ColumnarBatchSerializeResult(numRows, values) + val useOffheapBroadcastBuildRelation = + VeloxConfig.get.enableBroadcastBuildRelationInOffheap + new ColumnarBatchSerializeResult( + useOffheapBroadcastBuildRelation, + numRows, + values.toSeq.asJava) } else { ColumnarBatchSerializeResult.EMPTY } diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarCachedBatchSerializer.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarCachedBatchSerializer.scala index 80ff907e61e..a04e7d68fba 100644 --- a/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarCachedBatchSerializer.scala +++ b/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarCachedBatchSerializer.scala @@ -178,12 +178,7 @@ class ColumnarCachedBatchSerializer extends CachedBatchSerializer with Logging { BackendsApiManager.getBackendName, "ColumnarCachedBatchSerializer#serialize")) .serialize(ColumnarBatches.getNativeHandle(BackendsApiManager.getBackendName, batch)) - val bytes = - try { - unsafeBuffer.toByteArray - } finally { - unsafeBuffer.close() - } + val bytes = unsafeBuffer.toByteArray CachedColumnarBatch(batch.numRows(), bytes.length, bytes) } } diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeByteBufferArray.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeByteBufferArray.scala deleted file mode 100644 index 45ff118d6d0..00000000000 --- a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeByteBufferArray.scala +++ /dev/null @@ -1,138 +0,0 @@ -/* - * 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.execution.unsafe - -import org.apache.spark.annotation.Experimental -import org.apache.spark.internal.Logging -import org.apache.spark.memory.GlobalOffHeapMemory -import org.apache.spark.unsafe.Platform -import org.apache.spark.unsafe.array.LongArray -import org.apache.spark.unsafe.memory.MemoryAllocator - -/** - * Used to store broadcast variable off-heap memory for broadcast variable. The underlying data - * structure is a LongArray allocated in off-heap memory. - * - * @param arraySize - * underlying array[array[byte]]'s length - * @param byteBufferLengths - * underlying array[array[byte]] per byteBuffer length - * @param totalBytes - * all byteBuffer's length plus together - */ -// scalastyle:off no.finalize -@Experimental -case class UnsafeByteBufferArray(arraySize: Int, byteBufferLengths: Array[Int], totalBytes: Long) - extends Logging { - { - assert( - arraySize == byteBufferLengths.length, - "Unsafe buffer array size " + - "not equal to buffer lengths!") - assert(totalBytes >= 0, "Unsafe buffer array total bytes can't be negative!") - } - private val allocatedBytes = (totalBytes + 7) / 8 * 8 - - /** - * A single array to store all byteBufferArray's value, it's inited once when first time get - * accessed. - */ - private var longArray: LongArray = _ - - /** Index the start of each byteBuffer's offset to underlying LongArray's initial position. */ - private val byteBufferOffset = if (byteBufferLengths.isEmpty) { - new Array(0) - } else { - byteBufferLengths.init.scanLeft(0L)(_ + _) - } - - /** - * Put byteBuffer at specified array index. - * - * @param index - * index of the array. - * @param byteBuffer - * byteBuffer to put. - */ - def putByteBuffer(index: Int, byteBuffer: Array[Byte]): Unit = this.synchronized { - assert(index < arraySize) - assert(byteBuffer.length == byteBufferLengths(index)) - // first to allocate underlying long array - if (null == longArray && index == 0) { - GlobalOffHeapMemory.acquire(allocatedBytes) - longArray = new LongArray(MemoryAllocator.UNSAFE.allocate(allocatedBytes)) - } - - Platform.copyMemory( - byteBuffer, - Platform.BYTE_ARRAY_OFFSET, - longArray.getBaseObject, - longArray.getBaseOffset + byteBufferOffset(index), - byteBufferLengths(index)) - } - - /** - * Get byteBuffer at specified index. - * @param index - * @return - */ - def getByteBuffer(index: Int): Array[Byte] = { - assert(index < arraySize) - if (null == longArray) { - return new Array[Byte](0) - } - val bytes = new Array[Byte](byteBufferLengths(index)) - Platform.copyMemory( - longArray.getBaseObject, - longArray.getBaseOffset + byteBufferOffset(index), - bytes, - Platform.BYTE_ARRAY_OFFSET, - byteBufferLengths(index)) - bytes - } - - /** - * Get the byteBuffer memory address and length at specified index, usually used when read memory - * direct from offheap. - * - * @param index - * @return - */ - def getByteBufferOffsetAndLength(index: Int): (Long, Int) = { - assert(index < arraySize) - assert(longArray != null, "The broadcast data in offheap should not be null!") - val offset = longArray.getBaseOffset + byteBufferOffset(index) - val length = byteBufferLengths(index) - (offset, length) - } - - /** - * It's needed once the broadcast variable is garbage collected. Since now, we don't have an - * elegant way to free the underlying memory in offheap. - */ - override def finalize(): Unit = { - try { - if (longArray != null) { - longArray = null - GlobalOffHeapMemory.release(allocatedBytes) - } - } finally { - super.finalize() - } - } -} -// scalastyle:on no.finalize diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala index ebbd8226fd7..ba307415c50 100644 --- a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala +++ b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala @@ -47,10 +47,9 @@ import java.io.{Externalizable, ObjectInput, ObjectOutput} import scala.collection.JavaConverters.asScalaIteratorConverter object UnsafeColumnarBuildSideRelation { - // Keep constructors with BroadcastMode for compatibility def apply( output: Seq[Attribute], - batches: UnsafeByteBufferArray, + batches: Seq[UnsafeByteArray], mode: BroadcastMode): UnsafeColumnarBuildSideRelation = { val boundMode = mode match { case HashedRelationBroadcastMode(keys, isNullAware) => @@ -63,25 +62,6 @@ object UnsafeColumnarBuildSideRelation { } new UnsafeColumnarBuildSideRelation(output, batches, BroadcastModeUtils.toSafe(boundMode)) } - def apply( - output: Seq[Attribute], - byteBufferArray: Array[Array[Byte]], - mode: BroadcastMode): UnsafeColumnarBuildSideRelation = { - val boundMode = mode match { - case HashedRelationBroadcastMode(keys, isNullAware) => - // Bind each key to the build-side output so simple cols become BoundReference - val boundKeys: Seq[Expression] = - keys.map(k => BindReferences.bindReference(k, AttributeSeq(output))) - HashedRelationBroadcastMode(boundKeys, isNullAware) - case m => - m // IdentityBroadcastMode, etc. - } - new UnsafeColumnarBuildSideRelation( - output, - byteBufferArray, - BroadcastModeUtils.toSafe(boundMode) - ) - } } /** @@ -95,10 +75,10 @@ object UnsafeColumnarBuildSideRelation { * the broadcast mode. */ @Experimental -case class UnsafeColumnarBuildSideRelation( +class UnsafeColumnarBuildSideRelation( private var output: Seq[Attribute], - private var batches: UnsafeByteBufferArray, - var safeBroadcastMode: SafeBroadcastMode) + private var batches: Seq[UnsafeByteArray], + private var safeBroadcastMode: SafeBroadcastMode) extends BuildSideRelation with Externalizable with Logging @@ -118,93 +98,35 @@ case class UnsafeColumnarBuildSideRelation( /** needed for serialization. */ def this() = { - this(null, null.asInstanceOf[UnsafeByteBufferArray], null) + this(null, null, null) } - def this( - output: Seq[Attribute], - byteBufferArray: Array[Array[Byte]], - safeMode: SafeBroadcastMode - ) = { - this( - output, - UnsafeByteBufferArray( - byteBufferArray.length, - byteBufferArray.map(_.length), - byteBufferArray.map(_.length.toLong).sum - ), - safeMode - ) - val batchesSize = byteBufferArray.length - for (i <- 0 until batchesSize) { - // copy the bytes to off-heap memory. - batches.putByteBuffer(i, byteBufferArray(i)) - } + private[unsafe] def getBatches(): Seq[UnsafeByteArray] = { + batches } override def writeExternal(out: ObjectOutput): Unit = Utils.tryOrIOException { out.writeObject(output) out.writeObject(safeBroadcastMode) - out.writeInt(batches.arraySize) - out.writeObject(batches.byteBufferLengths) - out.writeLong(batches.totalBytes) - for (i <- 0 until batches.arraySize) { - val bytes = batches.getByteBuffer(i) - out.write(bytes) - } + out.writeObject(batches.toArray) } override def write(kryo: Kryo, out: Output): Unit = Utils.tryOrIOException { kryo.writeObject(out, output.toList) kryo.writeClassAndObject(out, safeBroadcastMode) - out.writeInt(batches.arraySize) - kryo.writeObject(out, batches.byteBufferLengths) - out.writeLong(batches.totalBytes) - for (i <- 0 until batches.arraySize) { - val bytes = batches.getByteBuffer(i) - out.write(bytes) - } + kryo.writeClassAndObject(out, batches.toArray) } override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException { output = in.readObject().asInstanceOf[Seq[Attribute]] safeBroadcastMode = in.readObject().asInstanceOf[SafeBroadcastMode] - val totalArraySize = in.readInt() - val byteBufferLengths = in.readObject().asInstanceOf[Array[Int]] - val totalBytes = in.readLong() - - // scalastyle:off - /** - * We use off-heap memory to reduce on-heap pressure Similar to - * https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala#L389-L410 - */ - // scalastyle:on - - batches = UnsafeByteBufferArray(totalArraySize, byteBufferLengths, totalBytes) - - for (i <- 0 until totalArraySize) { - val length = byteBufferLengths(i) - val tmpBuffer = new Array[Byte](length) - in.readFully(tmpBuffer) - batches.putByteBuffer(i, tmpBuffer) - } + batches = in.readObject().asInstanceOf[Array[UnsafeByteArray]].toSeq } override def read(kryo: Kryo, in: Input): Unit = Utils.tryOrIOException { output = kryo.readObject(in, classOf[List[_]]).asInstanceOf[Seq[Attribute]] safeBroadcastMode = kryo.readClassAndObject(in).asInstanceOf[SafeBroadcastMode] - val totalArraySize = in.readInt() - val byteBufferLengths = kryo.readObject(in, classOf[Array[Int]]) - val totalBytes = in.readLong() - - batches = UnsafeByteBufferArray(totalArraySize, byteBufferLengths, totalBytes) - - for (i <- 0 until totalArraySize) { - val length = byteBufferLengths(i) - val tmpBuffer = new Array[Byte](length) - in.read(tmpBuffer) - batches.putByteBuffer(i, tmpBuffer) - } + batches = kryo.readClassAndObject(in).asInstanceOf[Array[UnsafeByteArray]].toSeq } private def transformProjection: UnsafeProjection = safeBroadcastMode match { @@ -247,15 +169,17 @@ case class UnsafeColumnarBuildSideRelation( var batchId = 0 override def hasNext: Boolean = { - batchId < batches.arraySize + batchId < batches.size } override def next: ColumnarBatch = { - val (offset, length) = - batches.getByteBufferOffsetAndLength(batchId) + val unsafeByteArray = batches(batchId) batchId += 1 val handle = - jniWrapper.deserializeDirect(serializerHandle, offset, length) + jniWrapper.deserializeDirect( + serializerHandle, + unsafeByteArray.address(), + Math.toIntExact(unsafeByteArray.size())) ColumnarBatches.create(handle) } }) @@ -296,10 +220,10 @@ case class UnsafeColumnarBuildSideRelation( val jniWrapper = NativeColumnarToRowJniWrapper.create(runtime) val c2rId = jniWrapper.nativeColumnarToRowInit() var batchId = 0 - val iterator = if (batches.arraySize > 0) { + val iterator = if (batches.nonEmpty) { val res: Iterator[Iterator[InternalRow]] = new Iterator[Iterator[InternalRow]] { override def hasNext: Boolean = { - val itHasNext = batchId < batches.arraySize + val itHasNext = batchId < batches.size if (!itHasNext && !closed) { jniWrapper.nativeClose(c2rId) serializerJniWrapper.close(serializerHandle) @@ -309,10 +233,13 @@ case class UnsafeColumnarBuildSideRelation( } override def next(): Iterator[InternalRow] = { - val (offset, length) = batches.getByteBufferOffsetAndLength(batchId) + val unsafeByteArray = batches(batchId) batchId += 1 val batchHandle = - serializerJniWrapper.deserializeDirect(serializerHandle, offset, length) + serializerJniWrapper.deserializeDirect( + serializerHandle, + unsafeByteArray.address(), + Math.toIntExact(unsafeByteArray.size())) val batch = ColumnarBatches.create(batchHandle) if (batch.numRows == 0) { batch.close() @@ -370,7 +297,7 @@ case class UnsafeColumnarBuildSideRelation( override def estimatedSize: Long = { if (batches != null) { - batches.totalBytes + batches.map(_.size()).sum } else { 0L } diff --git a/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala b/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala index 903bb100b7b..6d0448fd84e 100644 --- a/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala +++ b/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala @@ -16,6 +16,8 @@ */ package org.apache.spark.sql.execution.unsafe +import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators + import org.apache.spark.{SparkConf, SparkEnv} import org.apache.spark.serializer.{JavaSerializer, KryoSerializer} import org.apache.spark.sql.catalyst.expressions.AttributeReference @@ -23,6 +25,9 @@ import org.apache.spark.sql.catalyst.plans.physical.IdentityBroadcastMode import org.apache.spark.sql.execution.joins.HashedRelationBroadcastMode import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.StringType +import org.apache.spark.unsafe.Platform + +import java.util class UnsafeColumnarBuildSideRelationTest extends SharedSparkSession { override protected def sparkConf: SparkConf = { @@ -33,28 +38,38 @@ class UnsafeColumnarBuildSideRelationTest extends SharedSparkSession { var unsafeRelWithIdentityMode: UnsafeColumnarBuildSideRelation = _ var unsafeRelWithHashMode: UnsafeColumnarBuildSideRelation = _ + var sampleBytes: Array[Array[Byte]] = _ + + private def toUnsafeByteArray(bytes: Array[Byte]): UnsafeByteArray = { + val buf = ArrowBufferAllocators.globalInstance().buffer(bytes.length) + buf.setBytes(0, bytes, 0, bytes.length); + new UnsafeByteArray(buf, bytes.length.toLong) + } + + private def toByteArray(unsafeByteArray: UnsafeByteArray): Array[Byte] = { + val byteArray = new Array[Byte](Math.toIntExact(unsafeByteArray.size())) + Platform.copyMemory( + null, + unsafeByteArray.address(), + byteArray, + Platform.BYTE_ARRAY_OFFSET, + byteArray.length) + byteArray + } override def beforeAll(): Unit = { super.beforeAll() val a = AttributeReference("a", StringType, nullable = false, null)() val output = Seq(a) - val totalArraySize = 1 - val perArraySize = new Array[Int](totalArraySize) - perArraySize(0) = 10 - val bytesArray = UnsafeByteBufferArray( - 1, - perArraySize, - 10 - ) - bytesArray.putByteBuffer(0, "1234567890".getBytes()) + sampleBytes = Array("12345".getBytes(), "7890".getBytes) unsafeRelWithIdentityMode = UnsafeColumnarBuildSideRelation( output, - bytesArray, + sampleBytes.map(a => toUnsafeByteArray(a)), IdentityBroadcastMode ) unsafeRelWithHashMode = UnsafeColumnarBuildSideRelation( output, - bytesArray, + sampleBytes.map(a => toUnsafeByteArray(a)), HashedRelationBroadcastMode(output, isNullAware = false) ) } @@ -68,12 +83,20 @@ class UnsafeColumnarBuildSideRelationTest extends SharedSparkSession { val obj = serializerInstance.deserialize[UnsafeColumnarBuildSideRelation](buffer) assert(obj != null) assert(obj.mode == IdentityBroadcastMode) + assert( + util.Arrays.deepEquals( + obj.getBatches().map(toByteArray).toArray[AnyRef], + sampleBytes.asInstanceOf[Array[AnyRef]])) // test unsafeRelWithHashMode val buffer2 = serializerInstance.serialize(unsafeRelWithHashMode) val obj2 = serializerInstance.deserialize[UnsafeColumnarBuildSideRelation](buffer2) assert(obj2 != null) assert(obj2.mode.isInstanceOf[HashedRelationBroadcastMode]) + assert( + util.Arrays.deepEquals( + obj2.getBatches().map(toByteArray).toArray[AnyRef], + sampleBytes.asInstanceOf[Array[AnyRef]])) } test("Kryo serialization") { @@ -85,12 +108,20 @@ class UnsafeColumnarBuildSideRelationTest extends SharedSparkSession { val obj = serializerInstance.deserialize[UnsafeColumnarBuildSideRelation](buffer) assert(obj != null) assert(obj.mode == IdentityBroadcastMode) + assert( + util.Arrays.deepEquals( + obj.getBatches().map(toByteArray).toArray[AnyRef], + sampleBytes.asInstanceOf[Array[AnyRef]])) // test unsafeRelWithHashMode val buffer2 = serializerInstance.serialize(unsafeRelWithHashMode) val obj2 = serializerInstance.deserialize[UnsafeColumnarBuildSideRelation](buffer2) assert(obj2 != null) assert(obj2.mode.isInstanceOf[HashedRelationBroadcastMode]) + assert( + util.Arrays.deepEquals( + obj2.getBatches().map(toByteArray).toArray[AnyRef], + sampleBytes.asInstanceOf[Array[AnyRef]])) } } diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc index f3b86f8ff6a..47de5061252 100644 --- a/cpp/core/jni/JniWrapper.cc +++ b/cpp/core/jni/JniWrapper.cc @@ -44,10 +44,10 @@ namespace { jclass byteArrayClass; -jclass unsafeByteBufferClass; -jmethodID unsafeByteBufferAllocate; -jmethodID unsafeByteBufferAddress; -jmethodID unsafeByteBufferSize; +jclass jniUnsafeByteBufferClass; +jmethodID jniUnsafeByteBufferAllocate; +jmethodID jniUnsafeByteBufferAddress; +jmethodID jniUnsafeByteBufferSize; jclass jniByteInputStreamClass; jmethodID jniByteInputStreamRead; @@ -250,11 +250,12 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) { byteArrayClass = createGlobalClassReferenceOrError(env, "[B"); - unsafeByteBufferClass = createGlobalClassReferenceOrError(env, "Lorg/apache/spark/unsafe/memory/UnsafeByteBuffer;"); - unsafeByteBufferAllocate = - env->GetStaticMethodID(unsafeByteBufferClass, "allocate", "(J)Lorg/apache/spark/unsafe/memory/UnsafeByteBuffer;"); - unsafeByteBufferAddress = env->GetMethodID(unsafeByteBufferClass, "address", "()J"); - unsafeByteBufferSize = env->GetMethodID(unsafeByteBufferClass, "size", "()J"); + jniUnsafeByteBufferClass = + createGlobalClassReferenceOrError(env, "Lorg/apache/spark/sql/execution/unsafe/JniUnsafeByteBuffer;"); + jniUnsafeByteBufferAllocate = env->GetStaticMethodID( + jniUnsafeByteBufferClass, "allocate", "(J)Lorg/apache/spark/sql/execution/unsafe/JniUnsafeByteBuffer;"); + jniUnsafeByteBufferAddress = env->GetMethodID(jniUnsafeByteBufferClass, "address", "()J"); + jniUnsafeByteBufferSize = env->GetMethodID(jniUnsafeByteBufferClass, "size", "()J"); jniByteInputStreamClass = createGlobalClassReferenceOrError(env, "Lorg/apache/gluten/vectorized/JniByteInputStream;"); jniByteInputStreamRead = getMethodIdOrError(env, jniByteInputStreamClass, "read", "(JJ)J"); @@ -298,7 +299,7 @@ void JNI_OnUnload(JavaVM* vm, void* reserved) { env->DeleteGlobalRef(splitResultClass); env->DeleteGlobalRef(nativeColumnarToRowInfoClass); env->DeleteGlobalRef(byteArrayClass); - env->DeleteGlobalRef(unsafeByteBufferClass); + env->DeleteGlobalRef(jniUnsafeByteBufferClass); env->DeleteGlobalRef(shuffleReaderMetricsClass); getJniErrorState()->close(); @@ -1164,9 +1165,9 @@ JNIEXPORT jobject JNICALL Java_org_apache_gluten_vectorized_ColumnarBatchSeriali auto serializer = ctx->createColumnarBatchSerializer(nullptr); serializer->append(batch); auto serializedSize = serializer->maxSerializedSize(); - auto byteBuffer = env->CallStaticObjectMethod(unsafeByteBufferClass, unsafeByteBufferAllocate, serializedSize); - auto byteBufferAddress = env->CallLongMethod(byteBuffer, unsafeByteBufferAddress); - auto byteBufferSize = env->CallLongMethod(byteBuffer, unsafeByteBufferSize); + auto byteBuffer = env->CallStaticObjectMethod(jniUnsafeByteBufferClass, jniUnsafeByteBufferAllocate, serializedSize); + auto byteBufferAddress = env->CallLongMethod(byteBuffer, jniUnsafeByteBufferAddress); + auto byteBufferSize = env->CallLongMethod(byteBuffer, jniUnsafeByteBufferSize); serializer->serializeTo(reinterpret_cast(byteBufferAddress), byteBufferSize); return byteBuffer; diff --git a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializeResult.java b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializeResult.java index 881048b052b..3aa26afd43a 100644 --- a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializeResult.java +++ b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializeResult.java @@ -16,26 +16,63 @@ */ package org.apache.gluten.vectorized; +import com.google.common.base.Preconditions; +import org.apache.spark.sql.execution.unsafe.JniUnsafeByteBuffer; +import org.apache.spark.sql.execution.unsafe.UnsafeByteArray; + import java.io.Serializable; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; public class ColumnarBatchSerializeResult implements Serializable { public static final ColumnarBatchSerializeResult EMPTY = - new ColumnarBatchSerializeResult(0, new byte[0][0]); - - private long numRows; + new ColumnarBatchSerializeResult(true, 0, Collections.emptyList()); - private byte[][] serialized; + private final boolean isOffHeap; + private final long numRows; + private final long sizeInBytes; + private final List onHeapData; + private final List offHeapData; - public ColumnarBatchSerializeResult(long numRows, byte[][] serialized) { + public ColumnarBatchSerializeResult( + boolean isOffHeap, long numRows, List serialized) { this.numRows = numRows; - this.serialized = serialized; + this.isOffHeap = isOffHeap; + if (isOffHeap) { + onHeapData = null; + offHeapData = + serialized.stream() + .map(JniUnsafeByteBuffer::toUnsafeByteArray) + .collect(Collectors.toList()); + sizeInBytes = offHeapData.stream().mapToInt(unsafe -> Math.toIntExact(unsafe.size())).sum(); + } else { + onHeapData = + serialized.stream().map(JniUnsafeByteBuffer::toByteArray).collect(Collectors.toList()); + offHeapData = null; + sizeInBytes = onHeapData.stream().mapToInt(bytes -> bytes.length).sum(); + } + } + + public boolean isOffHeap() { + return isOffHeap; } - public long getNumRows() { + public long numRows() { return numRows; } - public byte[][] getSerialized() { - return serialized; + public long sizeInBytes() { + return sizeInBytes; + } + + public List onHeapData() { + Preconditions.checkState(!isOffHeap); + return onHeapData; + } + + public List offHeapData() { + Preconditions.checkState(isOffHeap); + return offHeapData; } } diff --git a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializerJniWrapper.java b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializerJniWrapper.java index 00fe02c6293..909b5b411d1 100644 --- a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializerJniWrapper.java +++ b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializerJniWrapper.java @@ -19,7 +19,7 @@ import org.apache.gluten.runtime.Runtime; import org.apache.gluten.runtime.RuntimeAware; -import org.apache.spark.unsafe.memory.UnsafeByteBuffer; +import org.apache.spark.sql.execution.unsafe.JniUnsafeByteBuffer; public class ColumnarBatchSerializerJniWrapper implements RuntimeAware { private final Runtime runtime; @@ -37,7 +37,7 @@ public long rtHandle() { return runtime.getHandle(); } - public native UnsafeByteBuffer serialize(long handle); + public native JniUnsafeByteBuffer serialize(long handle); // Return the native ColumnarBatchSerializer handle public native long init(long cSchema); diff --git a/gluten-arrow/src/main/java/org/apache/spark/sql/execution/unsafe/JniUnsafeByteBuffer.java b/gluten-arrow/src/main/java/org/apache/spark/sql/execution/unsafe/JniUnsafeByteBuffer.java new file mode 100644 index 00000000000..86cbb5f7af2 --- /dev/null +++ b/gluten-arrow/src/main/java/org/apache/spark/sql/execution/unsafe/JniUnsafeByteBuffer.java @@ -0,0 +1,89 @@ +/* + * 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.execution.unsafe; + +import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.spark.unsafe.Platform; + +/** + * A temperate unsafe byte buffer implementation that is created and operated from C++ via JNI. The + * buffer has to be converted either to a Java on-heap byte array or to a Java off-heap unsafe byte + * array after Java code receives this object. + */ +public class JniUnsafeByteBuffer { + private ArrowBuf buffer; + private long size; + private boolean released = false; + + private JniUnsafeByteBuffer(ArrowBuf buffer, long size) { + this.buffer = buffer; + this.size = size; + } + + // Invoked by C++ code via JNI. + public static JniUnsafeByteBuffer allocate(long size) { + final ArrowBuf arrowBuf = ArrowBufferAllocators.globalInstance().buffer(size); + return new JniUnsafeByteBuffer(arrowBuf, size); + } + + // Invoked by C++ code via JNI. + public long address() { + ensureOpen(); + return buffer.memoryAddress(); + } + + // Invoked by C++ code via JNI. + public long size() { + ensureOpen(); + return size; + } + + private synchronized void ensureOpen() { + if (released) { + throw new IllegalStateException("Already released"); + } + } + + private synchronized void release() { + ensureOpen(); + buffer.close(); + released = true; + buffer = null; + size = 0; + } + + public synchronized byte[] toByteArray() { + ensureOpen(); + final byte[] values = new byte[Math.toIntExact(size)]; + Platform.copyMemory( + null, buffer.memoryAddress(), values, Platform.BYTE_ARRAY_OFFSET, values.length); + release(); + return values; + } + + public synchronized UnsafeByteArray toUnsafeByteArray() { + final UnsafeByteArray out; + ensureOpen(); + // We can safely release the buffer after UnsafeByteArray is constructed because it keeps + // its own reference to the buffer. + out = new UnsafeByteArray(buffer, size); + release(); + return out; + } +} diff --git a/gluten-arrow/src/main/java/org/apache/spark/sql/execution/unsafe/UnsafeByteArray.java b/gluten-arrow/src/main/java/org/apache/spark/sql/execution/unsafe/UnsafeByteArray.java new file mode 100644 index 00000000000..8a25beec21c --- /dev/null +++ b/gluten-arrow/src/main/java/org/apache/spark/sql/execution/unsafe/UnsafeByteArray.java @@ -0,0 +1,166 @@ +/* + * 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.execution.unsafe; + +import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators; + +import com.esotericsoftware.kryo.Kryo; +import com.esotericsoftware.kryo.KryoSerializable; +import com.esotericsoftware.kryo.io.Input; +import com.esotericsoftware.kryo.io.Output; +import org.apache.arrow.memory.ArrowBuf; + +import java.io.Externalizable; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; + +/** A serializable unsafe byte array. */ +public class UnsafeByteArray implements Externalizable, KryoSerializable { + private ArrowBuf buffer; + private long size; + + UnsafeByteArray(ArrowBuf buffer, long size) { + this.buffer = buffer; + this.buffer.getReferenceManager().retain(); + this.size = size; + } + + public UnsafeByteArray() {} + + public long address() { + return buffer.memoryAddress(); + } + + public long size() { + return size; + } + + public void release() { + if (buffer != null) { + buffer.close(); + buffer = null; + size = 0; + } + } + + // ------------ KryoSerializable ------------ + + @Override + public void write(Kryo kryo, Output output) { + // write length first + output.writeLong(size); + + // stream bytes out of ArrowBuf + final int chunkSize = 8 * 1024; + byte[] tmp = new byte[chunkSize]; + + long remaining = size; + int index = 0; + while (remaining > 0) { + int chunk = (int) Math.min(chunkSize, remaining); + buffer.getBytes(index, tmp, 0, chunk); + output.write(tmp, 0, chunk); + index += chunk; + remaining -= chunk; + } + } + + @Override + public void read(Kryo kryo, Input input) { + // read length + this.size = input.readLong(); + + if (size > Integer.MAX_VALUE) { + throw new IllegalArgumentException("UnsafeByteArray size too large: " + size); + } + + // allocate ArrowBuf + this.buffer = ArrowBufferAllocators.globalInstance().buffer((int) size); + + // stream bytes into ArrowBuf + final int chunkSize = 8 * 1024; + byte[] tmp = new byte[chunkSize]; + + long remaining = size; + int index = 0; + while (remaining > 0) { + int chunk = (int) Math.min(chunkSize, remaining); + input.readBytes(tmp, 0, chunk); + buffer.setBytes(index, tmp, 0, chunk); + index += chunk; + remaining -= chunk; + } + } + + // ------------ Externalizable ------------ + + @Override + public void writeExternal(ObjectOutput out) throws IOException { + // write length first + out.writeLong(size); + + final int chunkSize = 8 * 1024; + byte[] tmp = new byte[chunkSize]; + + long remaining = size; + int index = 0; + while (remaining > 0) { + int chunk = (int) Math.min(chunkSize, remaining); + buffer.getBytes(index, tmp, 0, chunk); + out.write(tmp, 0, chunk); + index += chunk; + remaining -= chunk; + } + } + + @Override + public void readExternal(ObjectInput in) throws IOException { + this.size = in.readLong(); + + if (size > Integer.MAX_VALUE) { + throw new IllegalArgumentException("UnsafeByteArray size too large: " + size); + } + + this.buffer = ArrowBufferAllocators.globalInstance().buffer((int) size); + + final int chunkSize = 8 * 1024; + byte[] tmp = new byte[chunkSize]; + + long remaining = size; + int index = 0; + while (remaining > 0) { + int chunk = (int) Math.min(chunkSize, remaining); + // ObjectInput extends DataInput, so we can use readFully + in.readFully(tmp, 0, chunk); + buffer.setBytes(index, tmp, 0, chunk); + index += chunk; + remaining -= chunk; + } + } + + /** + * It's needed once the broadcast variable is garbage collected. Since now, we don't have an + * elegant way to free the underlying memory in off-heap. + * + *

Since: https://github.com/apache/incubator-gluten/pull/8127. + */ + public void finalize() throws Throwable { + release(); + super.finalize(); + } +} diff --git a/gluten-arrow/src/main/java/org/apache/spark/unsafe/memory/UnsafeByteBuffer.java b/gluten-arrow/src/main/java/org/apache/spark/unsafe/memory/UnsafeByteBuffer.java deleted file mode 100644 index 1cd598ab6c7..00000000000 --- a/gluten-arrow/src/main/java/org/apache/spark/unsafe/memory/UnsafeByteBuffer.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * 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.unsafe.memory; - -import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators; - -import org.apache.arrow.memory.ArrowBuf; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.spark.task.TaskResources; -import org.apache.spark.unsafe.Platform; - -/** The API is for being called from C++ via JNI. */ -public class UnsafeByteBuffer { - private final ArrowBuf buffer; - private final long size; - - private UnsafeByteBuffer(ArrowBuf buffer, long size) { - this.buffer = buffer; - this.size = size; - } - - public static UnsafeByteBuffer allocate(long size) { - final BufferAllocator allocator; - if (TaskResources.inSparkTask()) { - allocator = ArrowBufferAllocators.contextInstance(UnsafeByteBuffer.class.getName()); - } else { - allocator = ArrowBufferAllocators.globalInstance(); - } - final ArrowBuf arrowBuf = allocator.buffer(size); - return new UnsafeByteBuffer(arrowBuf, size); - } - - public long address() { - return buffer.memoryAddress(); - } - - public long size() { - return size; - } - - public void close() { - buffer.close(); - } - - public byte[] toByteArray() { - final byte[] values = new byte[Math.toIntExact(size)]; - Platform.copyMemory( - null, buffer.memoryAddress(), values, Platform.BYTE_ARRAY_OFFSET, values.length); - return values; - } -}