westonpace commented on a change in pull request #12216: URL: https://github.com/apache/arrow/pull/12216#discussion_r794956943
########## File path: cpp/src/parquet/arrow/arrow_reader_writer_test.cc ########## @@ -3675,6 +3675,116 @@ TEST(TestArrowReaderAdHoc, LARGE_MEMORY_TEST(LargeStringColumn)) { AssertTablesEqual(*table, *batched_table, /*same_chunk_layout=*/false); } +class NoOverflowMemoryPool : public MemoryPool { + public: + explicit NoOverflowMemoryPool(MemoryPool* pool, int64_t padding); + ~NoOverflowMemoryPool() override = default; + + Status Allocate(int64_t size, uint8_t** out) override; + Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override; + + void Free(uint8_t* buffer, int64_t size) override; + + int64_t bytes_allocated() const override; + + int64_t max_memory() const override; + + std::string backend_name() const override; + + private: + MemoryPool* pool_; + int64_t padding_; + std::unordered_map<uint8_t**, int64_t> buffers_; + const uint8_t check_value_ = '\xff'; +}; + +NoOverflowMemoryPool::NoOverflowMemoryPool(MemoryPool* pool, int64_t padding) + : pool_(pool), padding_(padding) {} + +Status NoOverflowMemoryPool::Allocate(int64_t size, uint8_t** out) { + Status s = pool_->Allocate(size + padding_, out); + if (s.ok()) { + buffers_[out] = size; + + for (int i = size; i < size + padding_; ++i) { + (*out)[i] = check_value_; + } + } + return s; +} + +Status NoOverflowMemoryPool::Reallocate(int64_t old_size, int64_t new_size, + uint8_t** ptr) { + Status s = pool_->Reallocate(old_size + padding_, new_size + padding_, ptr); + if (s.ok()) { + buffers_[ptr] = new_size; + + for (int i = new_size; i < new_size + padding_; ++i) { + (*ptr)[i] = check_value_; + } + } + return s; +} + +void NoOverflowMemoryPool::Free(uint8_t* buffer, int64_t size) { + // DCHECK_EQ(size, buffers_[&buffer]); Review comment: I think you want `buffers_` to be `std::unordered_map<uint8_t*, int64_t> buffers_` instead of `uint8_t**` -- 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: github-unsubscr...@arrow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org