[FEA] Support unspill for SpillableHostBuffer

[binmahone]

Is your feature request related to a problem? Please describe.

Currently once a SpillableHostBuffer is spilled from memory to disk, all subsequent invocations of SpillableHostBuffer#getHostBuffer will read and deserialize from disk. It's very costly and won't be acceptable in cases where we will call the getHostBuffer multiple times.

One example would be the Kudo shuffle read concat case, let's asssume the read KudoTables are placed into a spillable state (by wrapping the HostMemoryBuffer in SpillableHostBuffer), then when doing the kudo concat, we will have to frequently and randomly call SpillableHostBuffer#getHostBuffer, since we know kudo concat adopts a random read visitor to read all the input KudoTables. It's a performance nightmere if we have to read from disk every time.

[abellina]

One example would be the Kudo shuffle read concat case, let's asssume the read KudoTables are placed into a spillable state (by wrapping the HostMemoryBuffer in SpillableHostBuffer), then when doing the kudo concat, we will have to frequently and randomly call SpillableHostBuffer#getHostBuffer, since we know kudo concat adopts a random read visitor to read all the input KudoTables. It's a performance nightmere if we have to read from disk every time.

I propose a change to the way the merge works then. I don't think unspill is the right way to solve the problem.

When you call materialize for a SpillableHostBuffer you get a HostMemoryBuffer that you must close after the operation is done. You are guaranteed the host memory. When we call mergeToTable and we are passing KudoTable those should just be all materialized on the host at once. In other words, expose a method in KudoTable that allows you materialize it, and hold a reference to the HostMemoryBuffer in KudoTable while you concat. When you are done, simply close the references.

[binmahone]

One example would be the Kudo shuffle read concat case, let's asssume the read KudoTables are placed into a spillable state (by wrapping the HostMemoryBuffer in SpillableHostBuffer), then when doing the kudo concat, we will have to frequently and randomly call SpillableHostBuffer#getHostBuffer, since we know kudo concat adopts a random read visitor to read all the input KudoTables. It's a performance nightmere if we have to read from disk every time.

I propose a change to the way the merge works then. I don't think unspill is the right way to solve the problem.

When you call materialize for a SpillableHostBuffer you get a HostMemoryBuffer that you must close after the operation is done. You are guaranteed the host memory. When we call mergeToTable and we are passing KudoTable those should just be all materialized on the host at once. In other words, expose a method in KudoTable that allows you materialize it, and hold a reference to the HostMemoryBuffer in KudoTable while you concat. When you are done, simply close the references.

Hi @abellina , I think you suggestion will also work, but I think somehow it breaks our assumption on "Spillable". I have drafted a PR (#12236) to fix #12215 accoding to your suggestion. But I have hightlighted the disadvantage of this apporach as in https://github.com/NVIDIA/spark-rapids/pull/12236/files#diff-f9b84ab98d00870a787854eb028b96af29e18eeffebbd71995c7371f144989aeR32.

As you can see #12236 does not require current issue (#12184) being solved

[binmahone]

Back to current issue (#12184), I figured unspill was something that will get done sooner or later, so I took the chance to get it done while I was at it. I also talked with @revans2 on this today, he said he'll have a dicussion with you on this. So I'm kind of waiting on your discussion results.

The unspill thing requires thorough considerations on mutlti-threading, and careful design of test cases, which takes significant amount of time. So if unspill is ultimately something we'll need, I suggest that we merge this PR.

It's also worth mentioning that once we have the unspill code checked in, we can use another pattern to "lock" the HostMemoryBuffer to avoid it being spilled again. Here's an example:

In the above snapshot, the code in green box are purely for "locking". By using this pattern we can avoid the problems described in https://github.com/NVIDIA/spark-rapids/pull/12236/files#diff-f9b84ab98d00870a787854eb028b96af29e18eeffebbd71995c7371f144989aeR32, because we don't need to add anything like guaranteeSpillable into KudoTable. However we do need to change the semantic of KudoTable.getBuffer(), after which caller has to remember to close the returned HostMemoryBuffer:

public class KudoTable implements AutoCloseable {
  private final KudoTableHeader header;
  private SpillableHostBuffer spillableHostBuffer;

   /**
   * Caller should close the buffer after use
   */
  public HostMemoryBuffer getBuffer() {
    return spillableHostBuffer.getHostBuffer(false);
  }

   /**
   * Caller should close the buffer after use
   */
  public HostMemoryBuffer getBuffer(boolean unspill) {
    return spillableHostBuffer.getHostBuffer(unspill);
  }
}