dongjoon-hyun commented on code in PR #59009:
URL: https://github.com/apache/spark/pull/59009#discussion_r4108469002
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -581,18 +609,28 @@ public Location lookup(Object keyBase, long keyOffset,
int keyLength) {
*
* This function always returns the same {@link Location} instance to avoid
object allocation.
* This function is not thread-safe.
+ *
+ * @throws IllegalStateException if the map is not ready for use
*/
public Location lookup(Object keyBase, long keyOffset, int keyLength, int
hash) {
safeLookup(keyBase, keyOffset, keyLength, loc, hash);
return loc;
}
+ private void ensureReady() {
Review Comment:
`Location.append()` is not guarded. With a `Location` obtained before a
failed `reset()` or `free()`, `append()` only has `assert (longArray != null)`,
so without `-ea` it acquires a new data page first and then hits an NPE at
`longArray.get(...)`, leaving that page in `dataPages`. Since `append()` runs
only for new keys, adding `ensureReady()` there should be cheap. Alternatively,
documenting that a stale `Location` is unsupported would be enough.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -1109,30 +1153,46 @@ public int getNumDataPages() {
}
/**
- * Returns the underline long[] of longArray.
+ * Returns the underlying long array.
+ *
+ * @throws IllegalStateException if the map is not ready for use
*/
public LongArray getArray() {
+ ensureReady();
assert(longArray != null);
return longArray;
}
/**
- * Reset this map to initialized state.
+ * Resets this map to initialized state. If the replacement pointer array
cannot be allocated,
+ * the map remains empty but cannot be used until a later call to this
method succeeds.
+ *
+ * @throws SparkOutOfMemoryError if the replacement pointer array cannot be
allocated
+ * @throws IllegalStateException if the map has been freed or destructive
iteration has begun
*/
public void reset() {
+ if (state == MapState.FREED || state == MapState.DESTRUCTIVE) {
+ throw new IllegalStateException("BytesToBytesMap cannot be reset in
state " + state);
+ }
+ state = MapState.RESET_FAILED;
updatePeakMemoryUsed();
+ // Put the map into its empty state before allocating so that a failed
allocation cannot leave
+ // stale page state. A caller may retry reset(), but cannot use the map
until reset succeeds.
Review Comment:
This (and the javadoc) still advertises retrying `reset()`. For the only
in-tree caller, `UnsafeKVExternalSorter`, a failed `reset()` also drops the
sorter holding the spilled groups, so a caller that retries `reset()` and keeps
aggregating would silently produce a result without the pre-spill groups. Could
we add a sentence that a caller retrying `reset()` is responsible for the
validity of any state kept outside the map (e.g., spilled data)?
##########
core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java:
##########
@@ -875,4 +875,178 @@ public void freeAfterFailedReset() {
}
}
+ @Test
+ public void operationsAfterFreeFail() {
+ memoryManager.limit(PAGE_SIZE_BYTES);
Review Comment:
nit. This `memoryManager.limit(PAGE_SIZE_BYTES)` does not seem to be needed
for this test.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -547,17 +566,24 @@ public Location next() {
*
* The returned iterator is NOT thread-safe. If the map is modified while
iterating over it,
* the behavior of the returned iterator is undefined.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration
has begun, or the map
+ * has been freed
*/
public MapIteratorWithKeyIndex iteratorWithKeyIndex() {
+ ensureReady();
return new MapIteratorWithKeyIndex();
}
/**
* The maximum number of allowed keys index.
*
* The value of allowed keys index is in the range of [0, maxNumKeysIndex -
1].
+ *
+ * @throws IllegalStateException if the map is not ready for use
Review Comment:
nit. The `@throws` wording is inconsistent: the iterator methods say "if a
reset failed, destructive iteration has begun, or the map has been freed",
while `lookup` / `safeLookup` / `getArray` / `maxNumKeysIndex` say "if the map
is not ready for use". Could we use one form?
##########
core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java:
##########
@@ -875,4 +875,178 @@ public void freeAfterFailedReset() {
}
}
+ @Test
+ public void operationsAfterFreeFail() {
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ map.free();
+ assertThrows(IllegalStateException.class, map::reset);
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertThrows(IllegalStateException.class, map::getArray);
+ assertThrows(IllegalStateException.class, map::maxNumKeysIndex);
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void configuredKeyOperationsRejectLookupAfterFree() {
+ BytesToBytesMap.KeyOperationsFactory keyOperationsFactory =
Review Comment:
nit. This anonymous `KeyOperations` duplicates the pattern of the existing
`caseInsensitiveWithConstantHash` fixture in this suite. It may be worth
sharing a fixture, although the existing one is test-local.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
