ORC-262: [C++] Support async I/O prefetch

c++/include/orc/OrcFile.hh

@@ -19,6 +19,7 @@
 #ifndef ORC_FILE_HH
 #define ORC_FILE_HH
 
+#include <future>
 #include <string>
 
 #include "orc/Reader.hh"
@@ -58,6 +59,18 @@ namespace orc {
      */
     virtual void read(void* buf, uint64_t length, uint64_t offset) = 0;
 
+    /**
+     * Read data asynchronously into the buffer. The buffer is allocated by the caller.
+     * @param buf the buffer to read into
+     * @param length the number of bytes to read.
+     * @param offset the position in the stream to read from.
+     * @return a future that will be set when the read is complete.
+     */
+    virtual std::future<void> readAsync(void* buf, uint64_t length, uint64_t offset) {
+      return std::async(std::launch::async,
+                        [this, buf, length, offset] { this->read(buf, length, offset); });
+    }
+
     /**
      * Get the name of the stream for error messages.
      */

c++/include/orc/Reader.hh

@@ -40,6 +40,17 @@ namespace orc {
   struct ReaderOptionsPrivate;
   struct RowReaderOptionsPrivate;
 
+  struct CacheOptions {
+    // The maximum distance in bytes between two consecutive
+    // ranges; beyond this value, ranges are not combined
+    uint64_t holeSizeLimit = 8192;
+
+    // The maximum size in bytes of a combined range; if
+    // combining two consecutive ranges would produce a range of a
+    // size greater than this, they are not combined
+    uint64_t rangeSizeLimit = 32 * 1024 * 1024;
+  };
+
   /**
    * Expose the reader metrics including the latency and
    * number of calls of the decompression/decoding/IO modules.
@@ -59,6 +70,8 @@ namespace orc {
     std::atomic<uint64_t> IOBlockingLatencyUs{0};
     std::atomic<uint64_t> SelectedRowGroupCount{0};
     std::atomic<uint64_t> EvaluatedRowGroupCount{0};
+    std::atomic<uint64_t> ReadRangeCacheHits{0};
+    std::atomic<uint64_t> ReadRangeCacheMisses{0};
   };
   ReaderMetrics* getDefaultReaderMetrics();
 
@@ -116,6 +129,11 @@ namespace orc {
      */
     ReaderOptions& setReaderMetrics(ReaderMetrics* metrics);
 
+    /**
+     * Set the cache options.
+     */
+    ReaderOptions& setCacheOptions(const CacheOptions& cacheOptions);
+
     /**
      * Set the location of the tail as defined by the logical length of the
      * file.
@@ -147,6 +165,11 @@ namespace orc {
      * Get the reader metrics.
      */
     ReaderMetrics* getReaderMetrics() const;
+
+    /**
+     * Set the cache options.
+     */
+    const CacheOptions& getCacheOptions() const;
   };
 
   /**
@@ -624,6 +647,23 @@ namespace orc {
      */
     virtual std::map<uint32_t, RowGroupIndex> getRowGroupIndex(
         uint32_t stripeIndex, const std::set<uint32_t>& included = {}) const = 0;
+
+    /**
+     * Trigger IO prefetch and cache the prefetched contents asynchronously.
+     * It is thread safe. Users should make sure requested stripes and columns
+     * are not overlapped, otherwise the overlapping part will be prefetched multiple time,
+     * which doesn't affect correctness but waste IO and memory resources.
+     * @param stripes the stripes to prefetch
+     * @param includeTypes the types to prefetch
+     */
+    virtual void preBuffer(const std::vector<uint32_t>& stripes,
+                           const std::list<uint64_t>& includeTypes) = 0;
+
+    /**
+     * Release cached entries whose right boundary is less than or equal to the given boundary.
+     * @param boundary the boundary value to release cache entries
+     */
+    virtual void releaseBuffer(uint64_t boundary) = 0;
   };
 
   /**

c++/src/CMakeLists.txt

@@ -150,6 +150,7 @@ set(SOURCE_FILES
   orc_proto.pb.h
   io/InputStream.cc
   io/OutputStream.cc
+  io/Cache.cc
   sargs/ExpressionTree.cc
   sargs/Literal.cc
   sargs/PredicateLeaf.cc

c++/src/Options.hh

@@ -23,6 +23,8 @@
 #include "orc/OrcFile.hh"
 #include "orc/Reader.hh"
 
+#include "io/Cache.hh"
+
 #include <limits>
 
 namespace orc {
@@ -43,6 +45,7 @@ namespace orc {
     MemoryPool* memoryPool;
     std::string serializedTail;
     ReaderMetrics* metrics;
+    CacheOptions cacheOptions;
 
     ReaderOptionsPrivate() {
       tailLocation = std::numeric_limits<uint64_t>::max();
@@ -122,6 +125,15 @@ namespace orc {
     return privateBits_->errorStream;
   }
 
+  ReaderOptions& ReaderOptions::setCacheOptions(const CacheOptions& cacheOptions) {
+    privateBits_->cacheOptions = cacheOptions;
+    return *this;
+  }
+
+  const CacheOptions& ReaderOptions::getCacheOptions() const {
+    return privateBits_->cacheOptions;
+  }
+
   /**
    * RowReaderOptions Implementation
    */

c++/src/Reader.cc

@@ -1523,6 +1523,94 @@ namespace orc {
     return ret;
   }
 
+  void ReaderImpl::releaseBuffer(uint64_t boundary) {
+    std::lock_guard<std::mutex> lock(contents_->readCacheMutex);
+
+    if (contents_->readCache) {
+      contents_->readCache->evictEntriesBefore(boundary);
+    }
+  }
+
+  void ReaderImpl::preBuffer(const std::vector<uint32_t>& stripes,
+                             const std::list<uint64_t>& includeTypes) {
+    std::vector<uint32_t> newStripes;
+    for (auto stripe : stripes) {
+      if (stripe < static_cast<uint32_t>(footer_->stripes_size())) newStripes.push_back(stripe);
+    }
+
+    std::list<uint64_t> newIncludeTypes;
+    for (auto type : includeTypes) {
+      if (type < static_cast<uint64_t>(footer_->types_size())) newIncludeTypes.push_back(type);
+    }
+
+    if (newStripes.empty() || newIncludeTypes.empty()) {
+      return;
+    }
+
+    orc::RowReaderOptions rowReaderOptions;
+    rowReaderOptions.includeTypes(newIncludeTypes);
+    ColumnSelector columnSelector(contents_.get());
+    std::vector<bool> selectedColumns;
+    columnSelector.updateSelected(selectedColumns, rowReaderOptions);
+
+    std::vector<ReadRange> ranges;
+    ranges.reserve(newIncludeTypes.size());
+    for (auto stripe : newStripes) {
+      // get stripe information
+      const auto& stripeInfo = footer_->stripes(stripe);
+      uint64_t stripeFooterStart =
+          stripeInfo.offset() + stripeInfo.index_length() + stripeInfo.data_length();
+      uint64_t stripeFooterLength = stripeInfo.footer_length();
+
+      // get stripe footer
+      std::unique_ptr<SeekableInputStream> pbStream = createDecompressor(
+          contents_->compression,
+          std::make_unique<SeekableFileInputStream>(contents_->stream.get(), stripeFooterStart,
+                                                    stripeFooterLength, *contents_->pool),
+          contents_->blockSize, *contents_->pool, contents_->readerMetrics);
+      proto::StripeFooter stripeFooter;
+      if (!stripeFooter.ParseFromZeroCopyStream(pbStream.get())) {
+        throw ParseError(std::string("bad StripeFooter from ") + pbStream->getName());
+      }
+
+      // traverse all streams in stripe footer, choose selected streams to prebuffer
+      uint64_t offset = stripeInfo.offset();
+      for (int i = 0; i < stripeFooter.streams_size(); i++) {
+        const proto::Stream& stream = stripeFooter.streams(i);
+        if (offset + stream.length() > stripeFooterStart) {
+          std::stringstream msg;
+          msg << "Malformed stream meta at stream index " << i << " in stripe " << stripe
+              << ": streamOffset=" << offset << ", streamLength=" << stream.length()
+              << ", stripeOffset=" << stripeInfo.offset()
+              << ", stripeIndexLength=" << stripeInfo.index_length()
+              << ", stripeDataLength=" << stripeInfo.data_length();
+          throw ParseError(msg.str());
+        }
+
+        if (stream.has_kind() && selectedColumns[stream.column()]) {
+          const auto& kind = stream.kind();
+          if (kind == proto::Stream_Kind_DATA || kind == proto::Stream_Kind_DICTIONARY_DATA ||
+              kind == proto::Stream_Kind_PRESENT || kind == proto::Stream_Kind_LENGTH ||
+              kind == proto::Stream_Kind_SECONDARY) {
+            ranges.emplace_back(offset, stream.length());
+          }
+        }
+
+        offset += stream.length();
+      }
+
+      {
+        std::lock_guard<std::mutex> lock(contents_->readCacheMutex);
+
+        if (!contents_->readCache) {
+          contents_->readCache = std::make_shared<ReadRangeCache>(
+              getStream(), options_.getCacheOptions(), contents_->pool, contents_->readerMetrics);
+        }
+        contents_->readCache->cache(std::move(ranges));
+      }
+    }
+  }
+
   RowReader::~RowReader() {
     // PASS
   }

c++/src/Reader.hh

@@ -26,6 +26,8 @@
 
 #include "ColumnReader.hh"
 #include "RLE.hh"
+#include "io/Cache.hh"
+
 #include "SchemaEvolution.hh"
 #include "TypeImpl.hh"
 #include "sargs/SargsApplier.hh"
@@ -70,6 +72,11 @@ namespace orc {
     bool isDecimalAsLong;
     std::unique_ptr<proto::Metadata> metadata;
     ReaderMetrics* readerMetrics;
+
+    // mutex to protect readCache_ from concurrent access
+    std::mutex readCacheMutex;
+    // cached io ranges. only valid when preBuffer is invoked.
+    std::shared_ptr<ReadRangeCache> readCache;
   };
 
   proto::StripeFooter getStripeFooter(const proto::StripeInformation& info,
@@ -245,6 +252,10 @@ namespace orc {
     const SchemaEvolution* getSchemaEvolution() const {
       return &schemaEvolution_;
     }
+
+    std::shared_ptr<ReadRangeCache> getReadCache() const {
+      return contents_->readCache;
+    }
   };
 
   class ReaderImpl : public Reader {
@@ -260,6 +271,7 @@ namespace orc {
     // footer
     proto::Footer* footer_;
     uint64_t numberOfStripes_;
+
     uint64_t getMemoryUse(int stripeIx, std::vector<bool>& selectedColumns);
 
     // internal methods
@@ -375,6 +387,10 @@ namespace orc {
     std::map<uint32_t, BloomFilterIndex> getBloomFilters(
         uint32_t stripeIndex, const std::set<uint32_t>& included) const override;
 
+    void preBuffer(const std::vector<uint32_t>& stripes,
+                   const std::list<uint64_t>& includeTypes) override;
+    void releaseBuffer(uint64_t boundary) override;
+
     std::map<uint32_t, RowGroupIndex> getRowGroupIndex(
         uint32_t stripeIndex, const std::set<uint32_t>& included) const override;
   };

c++/src/StripeStream.cc

@@ -19,6 +19,7 @@
 #include "StripeStream.hh"
 #include "RLE.hh"
 #include "Reader.hh"
+#include "io/Cache.hh"
 #include "orc/Exceptions.hh"
 
 #include "wrap/coded-stream-wrapper.h"
@@ -37,7 +38,8 @@ namespace orc {
         stripeStart_(stripeStart),
         input_(input),
         writerTimezone_(writerTimezone),
-        readerTimezone_(readerTimezone) {
+        readerTimezone_(readerTimezone),
+        readCache_(reader.getReadCache()) {
     // PASS
   }
 
@@ -89,7 +91,6 @@ namespace orc {
       if (stream.has_kind() && stream.kind() == kind &&
           stream.column() == static_cast<uint64_t>(columnId)) {
         uint64_t streamLength = stream.length();
-        uint64_t myBlock = shouldStream ? input_.getNaturalReadSize() : streamLength;
         if (offset + streamLength > dataEnd) {
           std::stringstream msg;
           msg << "Malformed stream meta at stream index " << i << " in stripe " << stripeIndex_
@@ -99,9 +100,23 @@ namespace orc {
               << ", stripeDataLength=" << stripeInfo_.data_length();
           throw ParseError(msg.str());
         }
-        return createDecompressor(reader_.getCompression(),
-                                  std::make_unique<SeekableFileInputStream>(
-                                      &input_, offset, stream.length(), *pool, myBlock),
+
+        BufferSlice slice;
+        if (readCache_) {
+          ReadRange range{offset, streamLength};
+          slice = readCache_->read(range);
+        }
+
+        uint64_t myBlock = shouldStream ? input_.getNaturalReadSize() : streamLength;
+        std::unique_ptr<SeekableInputStream> seekableInput;
+        if (slice.buffer) {
+          seekableInput = std::make_unique<SeekableArrayInputStream>(
+              slice.buffer->data() + slice.offset, slice.length);
+        } else {
+          seekableInput = std::make_unique<SeekableFileInputStream>(&input_, offset, streamLength,
+                                                                    *pool, myBlock);
+        }
+        return createDecompressor(reader_.getCompression(), std::move(seekableInput),
                                   reader_.getCompressionSize(), *pool,
                                   reader_.getFileContents().readerMetrics);
       }

c++/src/StripeStream.hh

@@ -30,6 +30,7 @@
 namespace orc {
 
   class RowReaderImpl;
+  class ReadRangeCache;
 
   /**
    * StripeStream Implementation
@@ -45,6 +46,7 @@ namespace orc {
     InputStream& input_;
     const Timezone& writerTimezone_;
     const Timezone& readerTimezone_;
+    std::shared_ptr<ReadRangeCache> readCache_;
 
    public:
     StripeStreamsImpl(const RowReaderImpl& reader, uint64_t index,