Make the file schema of written by parquet writer can be configurable (6074)

Author: rui-mo

velox/dwio/common/tests/CMakeLists.txt

@@ -64,6 +64,8 @@ target_link_libraries(
   velox_parse_parser
   velox_vector_test_lib
   velox_link_libs
+  parquet
+  arrow
   Folly::folly
   fmt::fmt
   lz4::lz4

velox/dwio/common/tests/E2EFilterTestBase.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <arrow/type.h>
 #include "velox/common/time/Timer.h"
 #include "velox/dwio/common/BufferedInput.h"
 #include "velox/dwio/common/FileSink.h"
@@ -170,7 +171,8 @@ class E2EFilterTestBase : public testing::Test {
   virtual void writeToMemory(
       const TypePtr& type,
       const std::vector<RowVectorPtr>& batches,
-      bool forRowGroupSkip) = 0;
+      bool forRowGroupSkip,
+      std::shared_ptr<arrow::Schema> schema = nullptr) = 0;
 
   virtual std::unique_ptr<dwio::common::Reader> makeReader(
       const dwio::common::ReaderOptions& opts,

velox/dwio/dwrf/test/E2EFilterTest.cpp

@@ -62,7 +62,8 @@ class E2EFilterTest : public E2EFilterTestBase {
   void writeToMemory(
       const TypePtr& type,
       const std::vector<RowVectorPtr>& batches,
-      bool forRowGroupSkip = false) override {
+      bool forRowGroupSkip = false,
+      std::shared_ptr<arrow::Schema> schema = nullptr) override {
     auto options = createWriterOptions(type);
     int32_t flushCounter = 0;
     // If we test row group skip, we have all the data in one stripe. For

velox/dwio/parquet/tests/reader/E2EFilterTest.cpp

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <arrow/type.h>
 #include "velox/dwio/common/tests/E2EFilterTestBase.h"
 #include "velox/dwio/parquet/reader/ParquetReader.h"
 #include "velox/dwio/parquet/writer/Writer.h"
@@ -69,6 +70,7 @@ class E2EFilterTest : public E2EFilterTestBase {
                 : (++flushCounter % flushEveryNBatches_ == 0);
           });
     };
+    options_.schema = schema;
 
     writer_ = std::make_unique<facebook::velox::parquet::Writer>(
         std::move(sink), options_);
@@ -615,6 +617,27 @@ TEST_F(E2EFilterTest, combineRowGroup) {
   EXPECT_EQ(parquetReader.numberOfRows(), 5);
 }
 
+TEST_F(E2EFilterTest, configurableWriteSchema) {
+  rowType_ = ROW({"c0"}, {INTEGER()});
+  std::vector<RowVectorPtr> batches;
+  for (int i = 0; i < 5; i++) {
+    batches.push_back(std::static_pointer_cast<RowVector>(
+        test::BatchMaker::createBatch(rowType_, 1, *leafPool_, nullptr, 0)));
+  }
+
+  auto int32 = arrow::field("int32", arrow::int32());
+  auto rbSchema = arrow::schema({int32});
+  writeToMemory(rowType_, batches, false, rbSchema);
+  std::string_view data(sinkPtr_->getData(), sinkPtr_->size());
+  dwio::common::ReaderOptions readerOpts{leafPool_.get()};
+  auto input = std::make_unique<BufferedInput>(
+      std::make_shared<InMemoryReadFile>(data), readerOpts.getMemoryPool());
+  auto reader = makeReader(readerOpts, std::move(input));
+  auto parquetReader = dynamic_cast<ParquetReader&>(*reader.get());
+  EXPECT_EQ(parquetReader.rowType()->containsChild("int32"), true);
+  EXPECT_EQ(parquetReader.rowType()->containsChild("c0"), false);
+}
+
 // Define main so that gflags get processed.
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);

velox/dwio/parquet/writer/Writer.cpp

@@ -141,6 +141,11 @@ Writer::Writer(
   } else {
     flushPolicy_ = std::make_unique<DefaultFlushPolicy>();
   }
+
+  if (options.schema) {
+    arrowContext_->schema = options.schema;
+  }
+
   arrowContext_->properties =
       getArrowParquetWriterOptions(options, flushPolicy_);
 }
@@ -218,11 +223,11 @@ void Writer::write(const VectorPtr& data) {
       auto recordBatch, arrow::ImportRecordBatch(&array, &schema));
   if (!arrowContext_->schema) {
     arrowContext_->schema = recordBatch->schema();
-    for (int colIdx = 0; colIdx < arrowContext_->schema->num_fields();
-         colIdx++) {
-      arrowContext_->stagingChunks.push_back(
-          std::vector<std::shared_ptr<arrow::Array>>());
-    }
+  }
+
+  for (int colIdx = 0; colIdx < arrowContext_->schema->num_fields(); colIdx++) {
+    arrowContext_->stagingChunks.push_back(
+        std::vector<std::shared_ptr<arrow::Array>>());
   }
 
   auto bytes = data->estimateFlatSize();

velox/dwio/parquet/writer/Writer.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <arrow/type.h>
 #include "velox/common/compression/Compression.h"
 #include "velox/dwio/common/DataBuffer.h"
 #include "velox/dwio/common/FileSink.h"
@@ -94,6 +95,7 @@ struct WriterOptions {
   // The default factory allows the writer to construct the default flush
   // policy with the configs in its ctor.
   std::function<std::unique_ptr<DefaultFlushPolicy>()> flushPolicyFactory;
+  std::shared_ptr<arrow::Schema> schema;
 };
 
 // Writes Velox vectors into  a DataSink using Arrow Parquet writer.