Skip to content

Commit c914a3c

Browse files
committed
address comments
1 parent 8b96561 commit c914a3c

3 files changed

Lines changed: 66 additions & 77 deletions

File tree

delta-lake/common/src/main/delta-33x-41x/scala/com/nvidia/spark/rapids/delta/common/RapidsDeletionVectors.scala

Lines changed: 47 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,21 @@ import org.apache.spark.sql.delta.storage.dv.HadoopFileSystemDVStore
3838
import org.apache.spark.sql.sources._
3939

4040
object RapidsDeletionVectors extends Logging {
41+
private def dvDescAndFilterType(
42+
dvDescriptorOpt: Option[String],
43+
filterTypeOpt: Option[RowIndexFilterType])
44+
: Option[(DeletionVectorDescriptor, RowIndexFilterType)] = {
45+
(dvDescriptorOpt, filterTypeOpt) match {
46+
case (Some(dvDescriptor), Some(filterType)) =>
47+
Some((DeletionVectorDescriptor.deserializeFromBase64(dvDescriptor), filterType))
48+
case (None, None) =>
49+
None
50+
case (Some(_), None) | (None, Some(_)) =>
51+
throw new IllegalStateException(
52+
"Both dvDescriptorOpt and filterTypeOpt must be defined together or both absent.")
53+
}
54+
}
55+
4156
/**
4257
* Translates the filter to use physical column names instead of logical column names.
4358
* This is needed when the column mapping mode is set to `NameMapping` or `IdMapping`
@@ -111,26 +126,22 @@ object RapidsDeletionVectors extends Logging {
111126
dvDescriptorOpt: Option[String],
112127
filterTypeOpt: Option[RowIndexFilterType],
113128
tablePath: String): HostMemoryBuffer = {
114-
if (dvDescriptorOpt.isDefined && filterTypeOpt.isDefined) {
115-
val dvDesc = DeletionVectorDescriptor.deserializeFromBase64(dvDescriptorOpt.get)
116-
117-
// The bitmap represents marked row indexes. The filter type determines whether those
118-
// rows are removed or retained.
119-
// See [[RowIndexFilterType]] for more details.
120-
filterTypeOpt.get match {
121-
case RowIndexFilterType.IF_CONTAINED | RowIndexFilterType.IF_NOT_CONTAINED =>
122-
val storedBitmap = RapidsDeletionVectorStoredBitmap(dvDesc, new Path(tablePath))
123-
storedBitmap.load(fileIO)
124-
case unexpectedFilterType => throw new IllegalStateException(
125-
s"Unexpected row index filter type for Deletion Vectors. " +
126-
s"Expected: ${RowIndexFilterType.IF_CONTAINED} or " +
127-
s"${RowIndexFilterType.IF_NOT_CONTAINED}; Actual: ${unexpectedFilterType}")
128-
}
129-
} else if (dvDescriptorOpt.isDefined || filterTypeOpt.isDefined) {
130-
throw new IllegalStateException(
131-
"Both dvDescriptorOpt and filterTypeOpt must be defined together or both absent.")
132-
} else {
133-
RapidsDeletionVectorStoredBitmap.serializedEmptyBitmap()
129+
dvDescAndFilterType(dvDescriptorOpt, filterTypeOpt) match {
130+
case Some((dvDesc, filterType)) =>
131+
// The bitmap represents marked row indexes. The filter type determines whether those
132+
// rows are removed or retained.
133+
// See [[RowIndexFilterType]] for more details.
134+
filterType match {
135+
case RowIndexFilterType.IF_CONTAINED | RowIndexFilterType.IF_NOT_CONTAINED =>
136+
val storedBitmap = RapidsDeletionVectorStoredBitmap(dvDesc, new Path(tablePath))
137+
storedBitmap.load(fileIO)
138+
case unexpectedFilterType => throw new IllegalStateException(
139+
s"Unexpected row index filter type for Deletion Vectors. " +
140+
s"Expected: ${RowIndexFilterType.IF_CONTAINED} or " +
141+
s"${RowIndexFilterType.IF_NOT_CONTAINED}; Actual: ${unexpectedFilterType}")
142+
}
143+
case None =>
144+
RapidsDeletionVectorStoredBitmap.serializedEmptyBitmap()
134145
}
135146
}
136147

@@ -151,26 +162,22 @@ object RapidsDeletionVectors extends Logging {
151162
dvDescriptorOpt: Option[String],
152163
filterTypeOpt: Option[RowIndexFilterType],
153164
tablePath: String): RoaringBitmapArray = {
154-
if (dvDescriptorOpt.isDefined && filterTypeOpt.isDefined) {
155-
val dvDesc = DeletionVectorDescriptor.deserializeFromBase64(dvDescriptorOpt.get)
156-
157-
// The bitmap represents marked row indexes. The filter type determines whether those
158-
// rows are removed or retained.
159-
// See [[RowIndexFilterType]] for more details.
160-
filterTypeOpt.get match {
161-
case RowIndexFilterType.IF_CONTAINED | RowIndexFilterType.IF_NOT_CONTAINED =>
162-
val dvStore = new HadoopFileSystemDVStore(conf)
163-
StoredBitmap.create(dvDesc, new Path(tablePath)).load(dvStore)
164-
case unexpectedFilterType => throw new IllegalStateException(
165-
s"Unexpected row index filter type for Deletion Vectors. " +
166-
s"Expected: ${RowIndexFilterType.IF_CONTAINED} or " +
167-
s"${RowIndexFilterType.IF_NOT_CONTAINED}; Actual: ${unexpectedFilterType}")
168-
}
169-
} else if (dvDescriptorOpt.isDefined || filterTypeOpt.isDefined) {
170-
throw new IllegalStateException(
171-
"Both dvDescriptorOpt and filterTypeOpt must be defined together or both absent.")
172-
} else {
173-
new RoaringBitmapArray()
165+
dvDescAndFilterType(dvDescriptorOpt, filterTypeOpt) match {
166+
case Some((dvDesc, filterType)) =>
167+
// The bitmap represents marked row indexes. The filter type determines whether those
168+
// rows are removed or retained.
169+
// See [[RowIndexFilterType]] for more details.
170+
filterType match {
171+
case RowIndexFilterType.IF_CONTAINED | RowIndexFilterType.IF_NOT_CONTAINED =>
172+
val dvStore = new HadoopFileSystemDVStore(conf)
173+
StoredBitmap.create(dvDesc, new Path(tablePath)).load(dvStore)
174+
case unexpectedFilterType => throw new IllegalStateException(
175+
s"Unexpected row index filter type for Deletion Vectors. " +
176+
s"Expected: ${RowIndexFilterType.IF_CONTAINED} or " +
177+
s"${RowIndexFilterType.IF_NOT_CONTAINED}; Actual: ${unexpectedFilterType}")
178+
}
179+
case None =>
180+
new RoaringBitmapArray()
174181
}
175182
}
176183

delta-lake/common/src/main/delta-33x-41x/scala/org/apache/spark/sql/delta/deletionvectors/RapidsDeletionVectorStore.scala

Lines changed: 16 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -83,15 +83,14 @@ private[deletionvectors] object RapidsInMemoryDeletionVectorStore {
8383
*/
8484
private sealed trait DeltaSerializedBitmapLoader {
8585
/**
86-
* Loads a bitmap payload and validates its trailing checksum. The CRC is initialized
87-
* with the magic number before this method is called.
86+
* Loads a bitmap payload and validates its trailing checksum. CRC validation is performed
87+
* when the `crc` parameter is provided. The bitmap is then converted to the "standard" roaring
88+
* bitmap serialization format, and returned as a HostMemoryBuffer.
8889
*/
89-
def loadAsStandardFormat(input: DataInputStream, size: Int, crc: CRC32): HostMemoryBuffer
90-
91-
/**
92-
* Loads an inline bitmap payload without CRC validation.
93-
*/
94-
def loadAsStandardFormat(input: DataInputStream, size: Int): HostMemoryBuffer
90+
def loadAsStandardFormat(
91+
input: DataInputStream,
92+
size: Int,
93+
crcOpt: Option[CRC32]): HostMemoryBuffer
9594
}
9695

9796
private object DeltaSerializedBitmapLoader {
@@ -133,18 +132,15 @@ private object DeltaSerializedBitmapLoader {
133132

134133
val remainingSize = size - DELTA_BITMAP_MAGIC_NUMBER_BYTE_SIZE
135134

136-
getFormatLoader(magicNumber).loadAsStandardFormat(input, remainingSize, crc)
135+
getFormatLoader(magicNumber).loadAsStandardFormat(input, remainingSize, Some(crc))
137136
}
138137

139138
def loadFromBytes(bytes: Array[Byte]): HostMemoryBuffer = {
140-
val bb = ByteBuffer.wrap(bytes)
141-
bb.order(ByteOrder.LITTLE_ENDIAN)
139+
val bb = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)
142140
val magicNumber = bb.getInt()
143141
val remainingSize = bb.remaining()
144-
withResource(new ByteBufferInputStream(bb)) { bais =>
145-
withResource(new DataInputStream(bais)) { in =>
146-
getFormatLoader(magicNumber).loadAsStandardFormat(in, remainingSize)
147-
}
142+
withResource(new DataInputStream(new ByteBufferInputStream(bb))) { in =>
143+
getFormatLoader(magicNumber).loadAsStandardFormat(in, remainingSize, None)
148144
}
149145
}
150146

@@ -162,17 +158,10 @@ private object DeltaSerializedBitmapLoader {
162158

163159
private object DeltaPortableFormatLoader extends DeltaSerializedBitmapLoader {
164160

165-
override def loadAsStandardFormat(input: DataInputStream, size: Int, crc: CRC32)
166-
: HostMemoryBuffer = loadAsStandardFormatImpl(input, size, Some(crc))
167-
168-
override def loadAsStandardFormat(input: DataInputStream, size: Int): HostMemoryBuffer = {
169-
loadAsStandardFormatImpl(input, size, None)
170-
}
171-
172-
private def loadAsStandardFormatImpl(
173-
input: DataInputStream,
174-
size: Int,
175-
crcOpt: Option[CRC32]): HostMemoryBuffer = {
161+
override def loadAsStandardFormat(
162+
input: DataInputStream,
163+
size: Int,
164+
crcOpt: Option[CRC32]): HostMemoryBuffer = {
176165
// The Delta portable format is identical to the standard portable format except for the
177166
// magic number at the beginning, which is already stripped at this point. Therefore,
178167
// we can directly load the remaining bytes into a HostMemoryBuffer and return it.
@@ -195,14 +184,7 @@ private object DeltaPortableFormatLoader extends DeltaSerializedBitmapLoader {
195184

196185
private object DeltaNativeFormatLoader extends DeltaSerializedBitmapLoader {
197186

198-
override def loadAsStandardFormat(input: DataInputStream, size: Int, crc: CRC32)
199-
: HostMemoryBuffer = loadAsStandardFormatImpl(input, size, Some(crc))
200-
201-
override def loadAsStandardFormat(input: DataInputStream, size: Int): HostMemoryBuffer = {
202-
loadAsStandardFormatImpl(input, size, None)
203-
}
204-
205-
private def loadAsStandardFormatImpl(
187+
override def loadAsStandardFormat(
206188
input: DataInputStream,
207189
size: Int,
208190
crcOpt: Option[CRC32]): HostMemoryBuffer = {

sql-plugin/src/main/scala/com/nvidia/spark/rapids/ByteBufferInputStream.scala

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ class ByteBufferInputStream(private var buffer: ByteBuffer)
2828
extends InputStream {
2929

3030
override def read(): Int = {
31-
if (buffer == null || buffer.remaining() == 0) {
31+
if (buffer == null || !buffer.hasRemaining()) {
3232
cleanUp()
3333
-1
3434
} else {
@@ -47,7 +47,7 @@ class ByteBufferInputStream(private var buffer: ByteBuffer)
4747
throw new IndexOutOfBoundsException
4848
} else if (length == 0) {
4949
0
50-
} else if (buffer == null || buffer.remaining() == 0) {
50+
} else if (buffer == null || !buffer.hasRemaining()) {
5151
cleanUp()
5252
-1
5353
} else {
@@ -61,7 +61,7 @@ class ByteBufferInputStream(private var buffer: ByteBuffer)
6161
if (buffer != null && bytes > 0) {
6262
val amountToSkip = math.min(bytes, buffer.remaining().toLong).toInt
6363
buffer.position(buffer.position() + amountToSkip)
64-
if (buffer.remaining() == 0) {
64+
if (!buffer.hasRemaining()) {
6565
cleanUp()
6666
}
6767
amountToSkip

0 commit comments

Comments
 (0)