[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 · Jan 3, 2024 · 957a022 · 957a022
1 parent 38de952
commit 957a022
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
@@ -163,11 +163,12 @@ std::shared_ptr<Datasource> VeloxRuntime::createDatasource(
     std::shared_ptr<arrow::Schema> schema) {
   static std::atomic_uint32_t id{0UL};
   auto veloxPool = getAggregateVeloxPool(memoryManager)->addAggregateChild("datasource." + std::to_string(id++));
-  // Pass a dedicate pool for S3 sink as can't share veloxPool
+  // Pass a dedicate pool 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
@@ -57,6 +57,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;
@@ -92,6 +97,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_;
 };