Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Update SpillableColumnarBatch to remove buffer from catalog on close #879

Merged
merged 3 commits into from
Sep 30, 2020
Merged

Update SpillableColumnarBatch to remove buffer from catalog on close #879

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
28bed7a
Update SpillableColumnarBatch to remove buffer from catalog on close
jlowe Sep 29, 2020
7ee2ddd
Add unit test
jlowe Sep 29, 2020
e1d1f67
Remove unused imports
jlowe Sep 29, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsBufferCatalog.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,12 @@ import java.util.concurrent.ConcurrentHashMap
import java.util.function.BiFunction

import ai.rapids.cudf.{DeviceMemoryBuffer, Rmm, Table}
import com.nvidia.spark.rapids.RapidsPluginImplicits.AutoCloseableArray
import com.nvidia.spark.rapids.StorageTier.StorageTier
import com.nvidia.spark.rapids.format.TableMeta

import org.apache.spark.{SparkConf, SparkEnv}
import org.apache.spark.internal.Logging
import org.apache.spark.sql.rapids.RapidsDiskBlockManager
import org.apache.spark.sql.vectorized.ColumnarBatch

/**
* Catalog for lookup of buffers by ID. The constructor is only visible for testing, generally
Expand Down Expand Up @@ -103,6 +101,9 @@ class RapidsBufferCatalog extends Logging {
buffer.free()
}
}

/** Return the number of buffers currently in the catalog. */
def numBuffers: Int = bufferMap.size()
}

object RapidsBufferCatalog extends Logging with Arm {
Expand Down Expand Up @@ -204,4 +205,7 @@ object RapidsBufferCatalog extends Logging with Arm {
* @return buffer that has been acquired
*/
def acquireBuffer(id: RapidsBufferId): RapidsBuffer = singleton.acquireBuffer(id)

/** Remove a buffer ID from the catalog and release the resources of the registered buffer. */
def removeBuffer(id: RapidsBufferId): Unit = singleton.removeBuffer(id)
}
4 changes: 1 addition & 3 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/SpillableColumnarBatch.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,9 +106,7 @@ class SpillableColumnarBatchImpl (id: TempSpillBufferId, rowCount: Int)
*/
override def close(): Unit = {
if (!closed) {
withResource(RapidsBufferCatalog.acquireBuffer(id)) { rapidsBuffer =>
rapidsBuffer.free()
}
RapidsBufferCatalog.removeBuffer(id)
closed = true
}
}
Expand Down
2 changes: 1 addition & 1 deletion sql-plugin/src/main/scala/org/apache/spark/sql/rapids/TempSpillBufferId.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ object TempSpillBufferId {
}
}

class TempSpillBufferId private(
case class TempSpillBufferId private(
override val tableId: Int,
bufferId: TempLocalBlockId) extends RapidsBufferId {

Expand Down