Handle buffered dwrf write exception to avoid server crash

facebookincubator · May 13, 2024 · 04d5bf2 · 04d5bf2
1 parent 49c3ebb
commit 04d5bf2
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/velox/dwio/dwrf/writer/DictionaryEncodingUtils.h b/velox/dwio/dwrf/writer/DictionaryEncodingUtils.h
@@ -80,6 +80,8 @@ class DictionaryEncodingUtils {
     uint32_t newIndex = 0;
     auto dictLengthWriter =
         createBufferedWriter<uint32_t>(pool, 64 * 1024, lengthFn);
+    auto errorGuard =
+        folly::makeGuard([&dictLengthWriter]() { dictLengthWriter.abort(); });
     for (uint32_t i = 0; i != numKeys; ++i) {
       auto origIndex = (sort ? sortedIndex[i] : i);
       if (!dropInfrequentKeys || shouldWriteKey(dictEncoder, origIndex)) {

diff --git a/velox/dwio/dwrf/writer/IntegerDictionaryEncoder.h b/velox/dwio/dwrf/writer/IntegerDictionaryEncoder.h
@@ -207,6 +207,7 @@ class IntegerDictionaryEncoder : public AbstractIntegerDictionaryEncoder {
     uint32_t newIndex = 0;
 
     auto dictWriter = createBufferedWriter<Integer>(writeBuffer, fn);
+    auto errorGuard = folly::makeGuard([&dictWriter]() { dictWriter.abort(); });
     for (uint32_t i = 0; i != numKeys; ++i) {
       auto origIndex = (sort ? sortedIndex[i] : i);
       if (!dropInfrequentKeys || shouldWriteKey(dictEncoder, origIndex)) {