From 0758a1d9d7d3ff1c987d7f6d646678101646b628 Mon Sep 17 00:00:00 2001 From: Shengxuan Liu Date: Tue, 28 Feb 2023 17:12:02 -0800 Subject: [PATCH] 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/reader/PageReader.cpp | 1 + velox/dwio/parquet/reader/ParquetReader.cpp | 1 + velox/dwio/parquet/thrift/CMakeLists.txt | 4 + .../parquet/thrift/ThriftBufferedTransport.h | 44 +++++++ .../parquet/thrift/ThriftStreamingTransport.h | 76 ++++++++++++ velox/dwio/parquet/thrift/ThriftTransport.h | 22 +--- .../dwio/parquet/thrift/tests/CMakeLists.txt | 20 ++++ .../thrift/tests/ThriftTransportTest.cpp | 112 ++++++++++++++++++ 8 files changed, 263 insertions(+), 17 deletions(-) create mode 100644 velox/dwio/parquet/thrift/ThriftBufferedTransport.h create mode 100644 velox/dwio/parquet/thrift/ThriftStreamingTransport.h create mode 100644 velox/dwio/parquet/thrift/tests/CMakeLists.txt create mode 100644 velox/dwio/parquet/thrift/tests/ThriftTransportTest.cpp diff --git a/velox/dwio/parquet/reader/PageReader.cpp b/velox/dwio/parquet/reader/PageReader.cpp index 5349afe941803..bb8e50d0b62ca 100644 --- a/velox/dwio/parquet/reader/PageReader.cpp +++ b/velox/dwio/parquet/reader/PageReader.cpp @@ -18,6 +18,7 @@ #include "velox/dwio/common/BufferUtil.h" #include "velox/dwio/common/ColumnVisitors.h" #include "velox/dwio/parquet/reader/NestedStructureDecoder.h" +#include "velox/dwio/parquet/thrift/ThriftBufferedTransport.h" #include "velox/dwio/parquet/thrift/ThriftTransport.h" #include "velox/vector/FlatVector.h" diff --git a/velox/dwio/parquet/reader/ParquetReader.cpp b/velox/dwio/parquet/reader/ParquetReader.cpp index 4485e1a9ba9d0..cea886c8cacde 100644 --- a/velox/dwio/parquet/reader/ParquetReader.cpp +++ b/velox/dwio/parquet/reader/ParquetReader.cpp @@ -19,6 +19,7 @@ #include "velox/dwio/common/MetricsLog.h" #include "velox/dwio/common/TypeUtils.h" #include "velox/dwio/parquet/reader/StructColumnReader.h" +#include "velox/dwio/parquet/thrift/ThriftBufferedTransport.h" #include "velox/dwio/parquet/thrift/ThriftTransport.h" namespace facebook::velox::parquet { diff --git a/velox/dwio/parquet/thrift/CMakeLists.txt b/velox/dwio/parquet/thrift/CMakeLists.txt index 78c2f3af1c8de..01e85ac2a8009 100644 --- a/velox/dwio/parquet/thrift/CMakeLists.txt +++ b/velox/dwio/parquet/thrift/CMakeLists.txt @@ -14,3 +14,7 @@ add_library(velox_dwio_parquet_thrift ParquetThriftTypes.cpp) target_link_libraries(velox_dwio_parquet_thrift arrow thrift) + +if(${VELOX_BUILD_TESTING}) + add_subdirectory(tests) +endif() diff --git a/velox/dwio/parquet/thrift/ThriftBufferedTransport.h b/velox/dwio/parquet/thrift/ThriftBufferedTransport.h new file mode 100644 index 0000000000000..00ba7bb3c068d --- /dev/null +++ b/velox/dwio/parquet/thrift/ThriftBufferedTransport.h @@ -0,0 +1,44 @@ +/* + * 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. + */ + +#pragma once + +#include +#include "velox/dwio/common/BufferedInput.h" +#include "velox/dwio/parquet/thrift/ThriftTransport.h" + +namespace facebook::velox::parquet::thrift { + +class ThriftBufferedTransport : public ThriftTransport { + public: + ThriftBufferedTransport(const void* inputBuf, uint64_t len) + : ThriftTransport(), + inputBuf_(reinterpret_cast(inputBuf)), + size_(len) {} + + uint32_t read(uint8_t* outputBuf, uint32_t len) { + DWIO_ENSURE(offset_ + len <= size_); + memcpy(outputBuf, inputBuf_ + offset_, len); + offset_ += len; + return len; + } + + private: + const uint8_t* inputBuf_; + const uint64_t size_; +}; + +} // namespace facebook::velox::parquet::thrift diff --git a/velox/dwio/parquet/thrift/ThriftStreamingTransport.h b/velox/dwio/parquet/thrift/ThriftStreamingTransport.h new file mode 100644 index 0000000000000..f96907ea50e37 --- /dev/null +++ b/velox/dwio/parquet/thrift/ThriftStreamingTransport.h @@ -0,0 +1,76 @@ +/* + * 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. + */ + +#pragma once + +#include +#include "velox/dwio/common/BufferedInput.h" +#include "velox/dwio/parquet/thrift/ThriftTransport.h" + +namespace facebook::velox::parquet::thrift { + +class ThriftStreamingTransport : public ThriftTransport { + public: + ThriftStreamingTransport( + std::shared_ptr inputStream, + const char* FOLLY_NULLABLE& bufferStart, + const char* FOLLY_NULLABLE& bufferEnd) + : ThriftTransport(), + 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 remainLen = len; + while (remainLen > 0) { + uint32_t size = bufferEnd_ - bufferStart_; + auto copySize = std::min(size, remainLen); + memcpy(outputBuf, bufferStart_, copySize); + remainLen -= copySize; + offset_ += copySize; + bufferStart_ += copySize; + outputBuf += copySize; + if (bufferStart_ == bufferEnd_) { + int32_t batchSize; + const void* bufferPointer; + if (!inputStream_->Next(&bufferPointer, &batchSize)) { + // The inputStream_ is fully consumed but reading does not + // finish. + if (remainLen > 0) { + VELOX_FAIL("Reading Exceeds inputStream_ Capacity"); + } + bufferStart_ = nullptr; + bufferEnd_ = nullptr; + } else { + bufferStart_ = static_cast(bufferPointer); + bufferEnd_ = bufferStart_ + batchSize; + } + } + } + return len; + } + + private: + std::shared_ptr inputStream_; + const char* FOLLY_NULLABLE& bufferStart_; + const char* FOLLY_NULLABLE& bufferEnd_; +}; + +} // namespace facebook::velox::parquet::thrift diff --git a/velox/dwio/parquet/thrift/ThriftTransport.h b/velox/dwio/parquet/thrift/ThriftTransport.h index 3e4bf769936fa..c0090073af53a 100644 --- a/velox/dwio/parquet/thrift/ThriftTransport.h +++ b/velox/dwio/parquet/thrift/ThriftTransport.h @@ -21,25 +21,13 @@ namespace facebook::velox::parquet::thrift { -class ThriftBufferedTransport - : public apache::thrift::transport::TVirtualTransport< - ThriftBufferedTransport> { +class ThriftTransport + : public apache::thrift::transport::TVirtualTransport { public: - ThriftBufferedTransport(const void* inputBuf, uint64_t len) - : inputBuf_(reinterpret_cast(inputBuf)), - size_(len), - offset_(0) {} + ThriftTransport() : offset_(0) {} + virtual uint32_t read(uint8_t* outputBuf, uint32_t len) = 0; - uint32_t read(uint8_t* outputBuf, uint32_t len) { - DWIO_ENSURE(offset_ + len <= size_); - memcpy(outputBuf, inputBuf_ + offset_, len); - offset_ += len; - return len; - } - - private: - const uint8_t* inputBuf_; - const uint64_t size_; + protected: uint64_t offset_; }; diff --git a/velox/dwio/parquet/thrift/tests/CMakeLists.txt b/velox/dwio/parquet/thrift/tests/CMakeLists.txt new file mode 100644 index 0000000000000..e3bb39100b702 --- /dev/null +++ b/velox/dwio/parquet/thrift/tests/CMakeLists.txt @@ -0,0 +1,20 @@ +# 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. +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error") + +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) diff --git a/velox/dwio/parquet/thrift/tests/ThriftTransportTest.cpp b/velox/dwio/parquet/thrift/tests/ThriftTransportTest.cpp new file mode 100644 index 0000000000000..4667fc07f39d9 --- /dev/null +++ b/velox/dwio/parquet/thrift/tests/ThriftTransportTest.cpp @@ -0,0 +1,112 @@ +/* + * 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 +#include +#include "velox/dwio/parquet/thrift/ThriftBufferedTransport.h" +#include "velox/dwio/parquet/thrift/ThriftStreamingTransport.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(i); + } + } + + void prepareThriftStreamingTransport() { + inputStream_ = std::make_shared( + 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(bufferPointer); + bufferEnd_ = bufferStart_ + batchSize_; + transport_ = std::make_shared( + inputStream_, bufferStart_, bufferEnd_); + } + + void prepareThriftBufferedTransport() { + transport_ = + std::make_shared(input_.data(), bufferSize_); + } + + static constexpr uint32_t bufferSize_ = 200; + static constexpr uint32_t batchSize_ = 20; + std::vector input_; + std::vector output_; + const char* FOLLY_NULLABLE bufferStart_{nullptr}; + const char* FOLLY_NULLABLE bufferEnd_{nullptr}; + std::shared_ptr inputStream_; + std::shared_ptr 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(); +}