[SPARK-40622][SQL][CORE]Remove the limitation that single task result must fit in 2GB #38064
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,47 @@ | ||
| /* | ||
| * 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.serializer | ||
|
|
||
| import java.nio.ByteBuffer | ||
|
|
||
| import scala.reflect.ClassTag | ||
|
|
||
| import org.apache.spark.internal.Logging | ||
| import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStream} | ||
|
|
||
| private[spark] object SerializerHelper extends Logging { | ||
|
Contributor
There was a problem hiding this comment. I was thinking if this is specific to the chunked byte buffer, would it be beneficial to add this to ChunkedByteBuffer as a static method, e.g.
Contributor
Author
There was a problem hiding this comment. Ideally, the |
||
| val CHUNK_BUFFER_SIZE: Int = 1024 * 1024 | ||
|
|
||
| def serializeToChunkedBuffer[T: ClassTag]( | ||
| serializerInstance: SerializerInstance, | ||
| t: T): ChunkedByteBuffer = { | ||
| val cbbos = new ChunkedByteBufferOutputStream(CHUNK_BUFFER_SIZE, ByteBuffer.allocate) | ||
| val out = serializerInstance.serializeStream(cbbos) | ||
| out.writeObject(t) | ||
| out.close() | ||
| cbbos.close() | ||
| cbbos.toChunkedByteBuffer | ||
| } | ||
liuzqt marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| def deserializeFromChunkedBuffer[T: ClassTag]( | ||
| serializerInstance: SerializerInstance, | ||
| bytes: ChunkedByteBuffer): T = { | ||
| val in = serializerInstance.deserializeStream(bytes.toInputStream()) | ||
| in.readObject() | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,7 +17,7 @@ | |
|
|
||
| package org.apache.spark.util.io | ||
|
|
||
| import java.io.{Externalizable, File, FileInputStream, InputStream, ObjectInput, ObjectOutput, OutputStream} | ||
| import java.nio.ByteBuffer | ||
| import java.nio.channels.WritableByteChannel | ||
|
|
||
|
|
@@ -42,8 +42,9 @@ import org.apache.spark.util.Utils | |
| * buffers may also be used elsewhere then the caller is responsible for copying | ||
| * them as needed. | ||
| */ | ||
| private[spark] class ChunkedByteBuffer(var chunks: Array[ByteBuffer]) extends Externalizable { | ||
|
Contributor
There was a problem hiding this comment. Can you confirm that these are copies of the byte buffers so there is no other reference that can be modified outside of the ChunkedByteBuffer based on your changes?
Contributor
Author
There was a problem hiding this comment. There is no physical bytes buffer copy happening here, so the But this PR can guarantee that the serialization won't modify any state of the ChunkedByteBuffer.
Contributor
There was a problem hiding this comment. I understand that, I was just asking if the references are maintained outside of this class, seems like it is not an issue.
Contributor
Author
There was a problem hiding this comment. Ahh got it. Yes in the case of this PR this should not be an issue since the chunk buffer is used within a pretty short life cycle during serialization/IO and reference ownership is clear. Thanks for the question. |
||
| require(chunks != null, "chunks must not be null") | ||
| require(!chunks.contains(null), "chunks must not contain null") | ||
|
Contributor
There was a problem hiding this comment. The next line should have implicitly checked for it - was this getting triggered ?
Contributor
Author
There was a problem hiding this comment. The It was triggered in some temp code snippet during my dev, but should not have any usage in the current code base(at least can not be found explicitly). That empty constructor might serve as default constructor for some serialization code I guess.......After all, the code in
Contributor
There was a problem hiding this comment. It is another iteration over |
||
| require(chunks.forall(_.position() == 0), "chunks' positions must be 0") | ||
|
|
||
| // Chunk size in bytes | ||
|
|
@@ -56,7 +57,13 @@ private[spark] class ChunkedByteBuffer(var chunks: Array[ByteBuffer]) { | |
| /** | ||
| * This size of this buffer, in bytes. | ||
| */ | ||
| var _size: Long = chunks.map(_.limit().asInstanceOf[Long]).sum | ||
liuzqt marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
|
|
||
| def size: Long = _size | ||
|
|
||
| def this() = { | ||
| this(Array.empty[ByteBuffer]) | ||
| } | ||
|
|
||
| def this(byteBuffer: ByteBuffer) = { | ||
| this(Array(byteBuffer)) | ||
|
|
@@ -84,6 +91,74 @@ private[spark] class ChunkedByteBuffer(var chunks: Array[ByteBuffer]) { | |
| } | ||
| } | ||
|
|
||
| /** | ||
| * write to ObjectOutput with zero copy if possible | ||
|
||
| */ | ||
| override def writeExternal(out: ObjectOutput): Unit = { | ||
| // we want to keep the chunks layout | ||
liuzqt marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| out.writeInt(chunks.length) | ||
| chunks.foreach(buffer => out.writeInt(buffer.limit())) | ||
| chunks.foreach(buffer => out.writeBoolean(buffer.isDirect)) | ||
| var buffer: Array[Byte] = null | ||
| val bufferLen = ChunkedByteBuffer.COPY_BUFFER_LEN | ||
|
|
||
| getChunks().foreach { chunk => { | ||
liuzqt marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| if (chunk.hasArray) { | ||
| // zero copy if the bytebuffer is backed by bytes array | ||
| out.write(chunk.array(), chunk.arrayOffset(), chunk.limit()) | ||
|
||
| } else { | ||
| // fallback to copy approach | ||
| if (buffer == null) { | ||
| buffer = new Array[Byte](bufferLen) | ||
| } | ||
| var bytesToRead = Math.min(chunk.remaining(), bufferLen) | ||
| while (bytesToRead > 0) { | ||
| chunk.get(buffer, 0, bytesToRead) | ||
| out.write(buffer, 0, bytesToRead) | ||
| bytesToRead = Math.min(chunk.remaining(), bufferLen) | ||
| } | ||
|
||
| } | ||
| }} | ||
| } | ||
|
|
||
| override def readExternal(in: ObjectInput): Unit = { | ||
| val chunksNum = in.readInt() | ||
| val indices = 0 until chunksNum | ||
| val chunksSize = indices.map(_ => in.readInt()) | ||
| val chunksDirect = indices.map(_ => in.readBoolean()) | ||
|
||
| val chunks = new Array[ByteBuffer](chunksNum) | ||
|
|
||
| val copyBufferLen = ChunkedByteBuffer.COPY_BUFFER_LEN | ||
| val copyBuffer: Array[Byte] = if (chunksDirect.exists(identity)) { | ||
| new Array[Byte](copyBufferLen) | ||
| } else { | ||
| null | ||
| } | ||
|
|
||
| indices.foreach { i => { | ||
liuzqt marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| val chunkSize = chunksSize(i) | ||
| chunks(i) = if (chunksDirect(i)) { | ||
| val buffer = ByteBuffer.allocateDirect(chunkSize) | ||
| var bytesRemaining = chunkSize | ||
| var bytesToRead = Math.min(copyBufferLen, bytesRemaining) | ||
| while (bytesRemaining > 0) { | ||
| bytesToRead = Math.min(copyBufferLen, bytesRemaining) | ||
| in.readFully(copyBuffer, 0, bytesToRead) | ||
| buffer.put(copyBuffer, 0, bytesToRead) | ||
| bytesRemaining -= bytesToRead | ||
| } | ||
| buffer.rewind() | ||
| buffer | ||
| } else { | ||
| val arr = new Array[Byte](chunkSize) | ||
| in.readFully(arr, 0, chunkSize) | ||
| ByteBuffer.wrap(arr) | ||
| } | ||
| }} | ||
| this.chunks = chunks | ||
| this._size = chunks.map(_.limit().toLong).sum | ||
| } | ||
|
|
||
| /** | ||
| * Wrap this in a custom "FileRegion" which allows us to transfer over 2 GB. | ||
| */ | ||
|
|
@@ -172,6 +247,8 @@ private[spark] class ChunkedByteBuffer(var chunks: Array[ByteBuffer]) { | |
|
|
||
| private[spark] object ChunkedByteBuffer { | ||
|
|
||
| val COPY_BUFFER_LEN: Int = 1024 * 1024 | ||
|
||
|
|
||
| def fromManagedBuffer(data: ManagedBuffer): ChunkedByteBuffer = { | ||
| data match { | ||
| case f: FileSegmentManagedBuffer => | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Unlike the earlier invocation of
serializeToChunkedBuffer(L#599) , here we have a good estimate of the size - something to leverage and minimize the cost ofserializeToChunkedBufferThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
do you mean we can explicitly choose the chunk size here to avoid too-small/to-large chunk?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, we do know
valueByteBuffer.size- and we know what the serialization overhead ofDirectTaskResultis.Will give us a much tighter bound on the size.