Refactor Thrift Transport for Parquet Metadata Access

This PR refactors the Thrift Transport for Parquet Metadata access. It uses
ThriftTransport as an interface and introduces ThriftBufferedTransport and
ThriftStreamingTransport.

ThriftStreamingTransport takes in a SeekableInputStream as input for Thrift
parsing. This can be used for Parquet Page Header parsing. This
optimization is able to reduce the deep copy in readPageHeader(). And it is
also the prerequisite to fix the incorrect page header length issue.

ThriftBufferedTransport takes in a consecutive memory space as input for
Thrift parsing. This can be used for Parquet Footer parsing, because the
footer is at the bottom of the file and we need to also take care of footer
length and PAR1, from the bottom to top.

velox/dwio/parquet/tests/CMakeLists.txt

@@ -24,6 +24,7 @@ set(TEST_LINK_LIBS
 
 add_subdirectory(duckdb_reader)
 add_subdirectory(reader)
+add_subdirectory(thrift)
 
 add_executable(velox_dwio_parquet_tpch_test ParquetTpchTest.cpp)
 add_test(

velox/dwio/parquet/tests/thrift/CMakeLists.txt

@@ -0,0 +1,19 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+add_executable(velox_dwio_parquet_thrift_test ThriftTransportTest.cpp)
+
+add_test(velox_dwio_parquet_thrift_test velox_dwio_parquet_thrift_test)
+target_link_libraries(velox_dwio_parquet_thrift_test arrow thrift
+                      ${VELOX_LINK_LIBS} gtest gtest_main)

velox/dwio/parquet/tests/thrift/ThriftTransportTest.cpp

@@ -0,0 +1,111 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "velox/dwio/parquet/thrift/ThriftTransport.h"
+#include <folly/init/Init.h>
+#include <gtest/gtest.h>
+
+using namespace facebook::velox;
+using namespace facebook::velox::dwio::common;
+using namespace facebook::velox::parquet::thrift;
+
+class ThriftTransportTest : public testing::Test {
+ protected:
+  void SetUp() override {
+    input_.resize(bufferSize_);
+    output_.resize(bufferSize_);
+    for (size_t i = 0; i < input_.size(); ++i) {
+      input_[i] = static_cast<uint8_t>(i);
+    }
+  }
+
+  void prepareThriftStreamingTransport() {
+    inputStream_ = std::make_shared<SeekableArrayInputStream>(
+        input_.data(), input_.size(), 20);
+    int32_t batchSize_;
+    const void* bufferPointer;
+    if (!inputStream_->Next(&bufferPointer, &batchSize_)) {
+      VELOX_CHECK(false, "Reading past end");
+    }
+    bufferStart_ = static_cast<const char*>(bufferPointer);
+    bufferEnd_ = bufferStart_ + batchSize_;
+    transport_ = std::make_shared<ThriftStreamingTransport>(
+        inputStream_.get(), bufferStart_, bufferEnd_);
+  }
+
+  void prepareThriftBufferedTransport() {
+    transport_ =
+        std::make_shared<ThriftBufferedTransport>(input_.data(), bufferSize_);
+  }
+
+  static constexpr uint32_t bufferSize_ = 200;
+  static constexpr uint32_t batchSize_ = 20;
+  std::vector<uint8_t> input_;
+  std::vector<uint8_t> output_;
+  const char* FOLLY_NULLABLE bufferStart_{nullptr};
+  const char* FOLLY_NULLABLE bufferEnd_{nullptr};
+  std::shared_ptr<SeekableInputStream> inputStream_;
+  std::shared_ptr<ThriftTransport> transport_;
+};
+
+TEST_F(ThriftTransportTest, streaming) {
+  prepareThriftStreamingTransport();
+  transport_->read(output_.data(), 10);
+  transport_->read(output_.data() + 10, 50);
+  transport_->read(output_.data() + 60, 140);
+
+  for (size_t i = 0; i < input_.size(); ++i) {
+    VELOX_CHECK_EQ(input_[i], output_[i]);
+  }
+}
+
+TEST_F(ThriftTransportTest, streamingOutOfBoundry) {
+  prepareThriftStreamingTransport();
+  transport_->read(output_.data(), 10);
+  transport_->read(output_.data() + 10, 50);
+  transport_->read(output_.data() + 60, 140);
+
+  // The whole inputStream_ is consumed.
+  EXPECT_ANY_THROW(transport_->read(output_.data() + bufferSize_, 1));
+}
+
+TEST_F(ThriftTransportTest, buffered) {
+  prepareThriftBufferedTransport();
+  transport_->read(output_.data(), 10);
+  transport_->read(output_.data() + 10, 50);
+  transport_->read(output_.data() + 60, 140);
+
+  for (size_t i = 0; i < input_.size(); ++i) {
+    VELOX_CHECK_EQ(input_[i], output_[i]);
+  }
+}
+
+TEST_F(ThriftTransportTest, bufferedOutOfBoundry) {
+  prepareThriftStreamingTransport();
+  transport_->read(output_.data(), 10);
+  transport_->read(output_.data() + 10, 50);
+  transport_->read(output_.data() + 60, 140);
+
+  // The whole inputStream_ is consumed.
+  EXPECT_ANY_THROW(transport_->read(output_.data() + bufferSize_, 1));
+}
+
+// Define main so that gflags get processed.
+int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+  folly::init(&argc, &argv, false);
+  return RUN_ALL_TESTS();
+}

velox/dwio/parquet/thrift/ThriftTransport.h

@@ -21,12 +21,62 @@
 
 namespace facebook::velox::parquet::thrift {
 
-class ThriftBufferedTransport
-    : public apache::thrift::transport::TVirtualTransport<
-          ThriftBufferedTransport> {
+class ThriftTransport
+    : public apache::thrift::transport::TVirtualTransport<ThriftTransport> {
+ public:
+  virtual uint32_t read(uint8_t* outputBuf, uint32_t len) = 0;
+  virtual ~ThriftTransport() = default;
+};
+
+class ThriftStreamingTransport : public ThriftTransport {
+ public:
+  ThriftStreamingTransport(
+      dwio::common::SeekableInputStream* inputStream,
+      const char*& bufferStart,
+      const char*& bufferEnd)
+      : inputStream_(inputStream),
+        bufferStart_(bufferStart),
+        bufferEnd_(bufferEnd) {
+
+    VELOX_CHECK_NOT_NULL(inputStream_);
+    VELOX_CHECK_NOT_NULL(bufferStart_);
+    VELOX_CHECK_NOT_NULL(bufferEnd_);
+  }
+
+  uint32_t read(uint8_t* outputBuf, uint32_t len) {
+    uint32_t bytesToRead = len;
+    while (bytesToRead > 0) {
+      if (bufferEnd_ == bufferStart_) {
+        int32_t size;
+        if (!inputStream_->Next(
+                reinterpret_cast<const void**>(&bufferStart_), &size)) {
+          VELOX_FAIL("Reading past the end of the stream");
+        }
+        bufferEnd_ = bufferStart_ + size;
+      }
+
+      uint32_t bytesToReadInBuffer =
+          std::min<uint32_t>(bufferEnd_ - bufferStart_, bytesToRead);
+      memcpy(outputBuf, bufferStart_, bytesToReadInBuffer);
+      bufferStart_ += bytesToReadInBuffer;
+      bytesToRead -= bytesToReadInBuffer;
+      outputBuf += bytesToReadInBuffer;
+    }
+
+    return len;
+  }
+
+ private:
+  dwio::common::SeekableInputStream* inputStream_;
+  const char*& bufferStart_;
+  const char*& bufferEnd_;
+};
+
+class ThriftBufferedTransport : public ThriftTransport {
  public:
   ThriftBufferedTransport(const void* inputBuf, uint64_t len)
-      : inputBuf_(reinterpret_cast<const uint8_t*>(inputBuf)),
+      : ThriftTransport(),
+        inputBuf_(reinterpret_cast<const uint8_t*>(inputBuf)),
         size_(len),
         offset_(0) {}