[VL] Add support to write parquet files to GCS

This change adds support to write parquet files to GCS. It is based on the support already present to write S3. Fixes #3976
apache · Dec 8, 2023 · 72162b5 · 72162b5
1 parent e9ea4e3
commit 72162b5
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 3 deletions.
diff --git a/cpp/velox/benchmarks/ParquetWriteBenchmark.cc b/cpp/velox/benchmarks/ParquetWriteBenchmark.cc
@@ -268,6 +268,7 @@ class GoogleBenchmarkVeloxParquetWriteCacheScanBenchmark : public GoogleBenchmar
           outputPath_ + "/" + fileName,
           veloxPool->addAggregateChild("writer_benchmark"),
           veloxPool->addLeafChild("s3_sink_pool"),
+          veloxPool->addLeafChild("gcs_sink_pool"),
           localSchema);
 
       veloxParquetDatasource->init(runtime->getConfMap());

diff --git a/cpp/velox/compute/VeloxRuntime.cc b/cpp/velox/compute/VeloxRuntime.cc
@@ -162,11 +162,12 @@ std::shared_ptr<Datasource> VeloxRuntime::createDatasource(
     MemoryManager* memoryManager,
     std::shared_ptr<arrow::Schema> schema) {
   auto veloxPool = getAggregateVeloxPool(memoryManager);
-  // Pass a dedicate pool for S3 sink as can't share veloxPool
+  // Pass dedicated pools for S3 and GCS sinks as can't share veloxPool
   // with parquet writer.
   auto s3SinkPool = getLeafVeloxPool(memoryManager);
+  auto gcsSinkPool = getLeafVeloxPool(memoryManager);
 
-  return std::make_shared<VeloxParquetDatasource>(filePath, veloxPool, s3SinkPool, schema);
+  return std::make_shared<VeloxParquetDatasource>(filePath, veloxPool, s3SinkPool, gcsSinkPool, schema);
 }
 
 std::shared_ptr<ShuffleReader> VeloxRuntime::createShuffleReader(

diff --git a/cpp/velox/operators/writer/VeloxParquetDatasource.cc b/cpp/velox/operators/writer/VeloxParquetDatasource.cc
@@ -56,6 +56,16 @@ void VeloxParquetDatasource::init(const std::unordered_map<std::string, std::str
 #else
     throw std::runtime_error(
         "The write path is S3 path but the S3 haven't been enabled when writing parquet data in velox runtime!");
+#endif
+  } else if (strncmp(filePath_.c_str(), "gs:", 3) == 0) {
+#ifdef ENABLE_GCS
+    auto fileSystem = getFileSystem(filePath_, nullptr);
+    auto* gcsFileSystem = dynamic_cast<filesystems::GCSFileSystem*>(fileSystem.get());
+    sink_ = std::make_unique<dwio::common::WriteFileSink>(
+        gcsFileSystem->openFileForWrite(filePath_, {{}, gcsSinkPool_.get()}), filePath_);
+#else
+    throw std::runtime_error(
+        "The write path is GCS path but the GCS haven't been enabled when writing parquet data in velox runtime!");
 #endif
   } else if (strncmp(filePath_.c_str(), "hdfs:", 5) == 0) {
 #ifdef ENABLE_HDFS

diff --git a/cpp/velox/operators/writer/VeloxParquetDatasource.h b/cpp/velox/operators/writer/VeloxParquetDatasource.h
@@ -36,6 +36,9 @@
 #include "velox/connectors/hive/storage_adapters/s3fs/S3FileSystem.h"
 #include "velox/connectors/hive/storage_adapters/s3fs/S3Util.h"
 #endif
+#ifdef ENABLE_GCS
+#include "velox/connectors/hive/storage_adapters/gcs/GCSFileSystem.h"
+#endif
 #ifdef ENABLE_HDFS
 #include "velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.h"
 #include "velox/connectors/hive/storage_adapters/hdfs/HdfsUtil.h"
@@ -54,12 +57,14 @@ class VeloxParquetDatasource final : public Datasource {
       const std::string& filePath,
       std::shared_ptr<facebook::velox::memory::MemoryPool> veloxPool,
       std::shared_ptr<facebook::velox::memory::MemoryPool> s3SinkPool,
+      std::shared_ptr<facebook::velox::memory::MemoryPool> gcsSinkPool,
       std::shared_ptr<arrow::Schema> schema)
       : Datasource(filePath, schema),
         filePath_(filePath),
         schema_(schema),
         pool_(std::move(veloxPool)),
-        s3SinkPool_(std::move(s3SinkPool)) {}
+        s3SinkPool_(std::move(s3SinkPool)),
+        gcsSinkPool_(std::move(gcsSinkPool)) {}
 
   void init(const std::unordered_map<std::string, std::string>& sparkConfs) override;
   void inspectSchema(struct ArrowSchema* out) override;
@@ -79,6 +84,7 @@ class VeloxParquetDatasource final : public Datasource {
   std::shared_ptr<facebook::velox::parquet::Writer> parquetWriter_;
   std::shared_ptr<facebook::velox::memory::MemoryPool> pool_;
   std::shared_ptr<facebook::velox::memory::MemoryPool> s3SinkPool_;
+  std::shared_ptr<facebook::velox::memory::MemoryPool> gcsSinkPool_;
   std::unique_ptr<facebook::velox::dwio::common::FileSink> sink_;
 };