From 04d5bf2c3af9e4140196b3350eee604b70287bbf Mon Sep 17 00:00:00 2001
From: xiaoxmeng <xiaoxmeng@fb.com>
Date: Sun, 12 May 2024 21:20:41 -0700
Subject: [PATCH] Handle buffered dwrf write exception to avoid server crash

---
 velox/dwio/dwrf/writer/DictionaryEncodingUtils.h  | 2 ++
 velox/dwio/dwrf/writer/IntegerDictionaryEncoder.h | 1 +
 2 files changed, 3 insertions(+)

diff --git a/velox/dwio/dwrf/writer/DictionaryEncodingUtils.h b/velox/dwio/dwrf/writer/DictionaryEncodingUtils.h
index 70ecd3a9aa7c..912fc4dc61a9 100644
--- 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
index 69c3a18694df..4bc772d2ac00 100644
--- 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)) {