refactor(parquet): Use velox parquet reader in FileSerializeTest

facebookincubator · Dec 17, 2024 · 4abc922 · 4abc922
1 parent f1622ab
commit 4abc922
Show file tree

Hide file tree

Showing 7 changed files with 274 additions and 94 deletions.
diff --git a/velox/dwio/parquet/reader/Metadata.cpp b/velox/dwio/parquet/reader/Metadata.cpp
@@ -210,6 +210,11 @@ bool ColumnChunkMetaDataPtr::hasDictionaryPageOffset() const {
       thriftColumnChunkPtr(ptr_)->meta_data.__isset.dictionary_page_offset;
 }
 
+bool ColumnChunkMetaDataPtr::hasIndexPage() const {
+  return hasMetadata() &&
+      thriftColumnChunkPtr(ptr_)->meta_data.__isset.index_page_offset;
+}
+
 std::unique_ptr<dwio::common::ColumnStatistics>
 ColumnChunkMetaDataPtr::getColumnStatistics(
     const TypePtr type,

diff --git a/velox/dwio/parquet/reader/Metadata.h b/velox/dwio/parquet/reader/Metadata.h
@@ -37,6 +37,8 @@ class ColumnChunkMetaDataPtr {
   /// Check the presence of the dictionary page offset in ColumnChunk metadata.
   bool hasDictionaryPageOffset() const;
 
+  bool hasIndexPage() const;
+
   /// Return the ColumnChunk statistics.
   std::unique_ptr<dwio::common::ColumnStatistics> getColumnStatistics(
       const TypePtr type,

diff --git a/velox/dwio/parquet/reader/ParquetReader.cpp b/velox/dwio/parquet/reader/ParquetReader.cpp
@@ -18,6 +18,7 @@
 
 #include <thrift/protocol/TCompactProtocol.h> //@manual
 
+#include <arrow/io/api.h>
 #include "velox/dwio/parquet/reader/ParquetColumnReader.h"
 #include "velox/dwio/parquet/reader/StructColumnReader.h"
 #include "velox/dwio/parquet/thrift/ThriftTransport.h"
@@ -843,10 +844,12 @@ std::shared_ptr<const RowType> ReaderBase::createRowType(
   std::vector<TypePtr> childTypes;
   for (auto& child : children) {
     auto childName = static_cast<const ParquetTypeWithId&>(*child).name_;
+    auto childt = static_cast<const ParquetTypeWithId&>(*child).parquetType_;
     if (fileColumnNamesReadAsLowerCase) {
       folly::toLowerAscii(childName);
     }
     childNames.push_back(std::move(childName));
+    // childTypes.push_back(childt);
     childTypes.push_back(child->type());
   }
   return TypeFactory<TypeKind::ROW>::create(
@@ -1156,4 +1159,7 @@ FileMetaDataPtr ParquetReader::fileMetaData() const {
   return readerBase_->fileMetaData();
 }
 
+const thrift::FileMetaData& ParquetReader::thriftFileMetaData() const {
+  return readerBase_->thriftFileMetaData();
+}
 } // namespace facebook::velox::parquet
diff --git a/velox/dwio/parquet/reader/ParquetReader.h b/velox/dwio/parquet/reader/ParquetReader.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <arrow/io/api.h>
 #include "velox/dwio/common/Reader.h"
 #include "velox/dwio/common/ReaderFactory.h"
 #include "velox/dwio/parquet/reader/Metadata.h"
@@ -106,6 +107,8 @@ class ParquetReader : public dwio::common::Reader {
 
   FileMetaDataPtr fileMetaData() const;
 
+  const thrift::FileMetaData& thriftFileMetaData() const;
+
  private:
   std::shared_ptr<ReaderBase> readerBase_;
 };
@@ -125,4 +128,8 @@ void registerParquetReaderFactory();
 
 void unregisterParquetReaderFactory();
 
+struct ColumnInfo {
+  int64_t col_start;
+  int64_t col_length;
+};
 } // namespace facebook::velox::parquet
diff --git a/velox/dwio/parquet/tests/writer/CMakeLists.txt b/velox/dwio/parquet/tests/writer/CMakeLists.txt
@@ -48,3 +48,21 @@ target_link_libraries(
   ${TEST_LINK_LIBS}
   GTest::gtest
   fmt::fmt)
+
+add_executable(velox_dwio_parquet_writer_test FileSerializeTest.cpp)
+
+add_test(
+  NAME velox_dwio_parquet_writer_test
+  COMMAND velox_dwio_parquet_writer_test
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+
+target_link_libraries(
+  velox_dwio_parquet_writer_test
+  velox_dwio_arrow_parquet_writer_test_lib
+  GTest::gmock
+  GTest::gtest
+  GTest::gtest_main
+  arrow
+  arrow_testing
+  velox_dwio_native_parquet_reader
+  velox_temp_path)