make kudo shuffle read retryable and spillable #12236
Conversation
Signed-off-by: Hongbin Ma (Mahone) <mahongbin@apache.org>
| import com.nvidia.spark.rapids.jni.kudo.KudoTable; | ||
| import com.nvidia.spark.rapids.jni.kudo.KudoTableHeader; | ||
|
|
||
| public class SpillableKudoTable extends KudoTable { |
There was a problem hiding this comment.
if this was in scala, it might make more sense, as it's only used from scala.
There was a problem hiding this comment.
For me I would prefer to have SpillableKudoTable not extend KudoTable.
Conceptually I want to say that I have a list of SpillableKudoTables, and then when I want to concat and serialize them into a Table I convert all of them into regular KudoTables which guarantees that they are all resident in memory, do the concat operation, and then close them when I am done.
This implies that a SpillableKudoTable is a KudoTable, which we can make it work, but as @abellina pointed out it makes the code much more complicated. If you have SpillableKudoTable produce a KudoTable, then we don't need guaranteeSpillable. The act of getting the KudoTable out implicitly increments the reference count and make it not spillable until it goes out of scope.
| import com.nvidia.spark.rapids.jni.kudo.KudoTableHeader; | ||
|
|
||
| public class SpillableKudoTable extends KudoTable { | ||
| private SpillableHostBuffer shb; |
There was a problem hiding this comment.
really this should just be a SpillableHostBufferHandle, but for now it probably fits better with the code.
|
|
||
| KudoHostMergeResultWrapper(result) | ||
| try { | ||
| for (skt <- columns.map(_.spillableKudoTable)) { |
There was a problem hiding this comment.
could a SpillableKudoTable make a KudoTable that isn't spillable? I don't know if SpillableKudoTable has to subclass KudoTable, but I feel this would clean this interface a bit.
So you say skts.safeMap(_.makeKudoTable). KudoTable is already closeable, so when we close the sequence of KudoTable all HostMemoryBuffer references are closed and we are now spillable again.
We then later skts.safeClose() to close the spillable handles.
This PR fixes #12215 in the way suggested by @abellina in #12184 (comment). It requires NVIDIA/spark-rapids-jni#2991 being merged first