Merge pull request #42 from mhaseeb123/fea/cudf-parquet-connector

Add a new `ParquetConnector` and `ParquetDataSource` to read Parquet tables via cuDF backend

Author: mhaseeb123

velox/experimental/cudf/CMakeLists.txt

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 add_subdirectory(exec)
+add_subdirectory(connectors)
 add_subdirectory(vector)
 
 if(VELOX_BUILD_TESTING)

velox/experimental/cudf/connectors/CMakeLists.txt

@@ -0,0 +1,15 @@
+# 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.
+
+add_subdirectory(parquet)

velox/experimental/cudf/connectors/parquet/CMakeLists.txt

@@ -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.
+
+add_library(velox_cudf_parquet_reader_config ParquetReaderConfig.cpp)
+
+set_target_properties(
+  velox_cudf_parquet_reader_config
+  PROPERTIES CUDA_ARCHITECTURES native)
+
+target_link_libraries(
+  velox_cudf_parquet_reader_config velox_core velox_exception cudf::cudf)
+
+add_library(
+  velox_cudf_parquet_connector OBJECT
+  ParquetReaderConfig.cpp
+  ParquetConnector.cpp
+  ParquetConnectorSplit.cpp
+  ParquetDataSource.cpp
+  ParquetTableHandle.cpp)
+
+set_target_properties(
+  velox_cudf_parquet_connector
+  PROPERTIES CUDA_ARCHITECTURES native)
+
+target_link_libraries(
+  velox_cudf_parquet_connector
+  PRIVATE
+    cudf::cudf
+    velox_cudf_exec
+    velox_common_io
+    velox_connector
+    velox_type_tz
+    velox_gcs)

velox/experimental/cudf/connectors/parquet/ParquetConnector.cpp

@@ -0,0 +1,56 @@
+/*
+ * 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 "velox/experimental/cudf/connectors/parquet/ParquetConnector.h"
+#include "velox/experimental/cudf/connectors/parquet/ParquetDataSource.h"
+
+namespace facebook::velox::cudf_velox::connector::parquet {
+
+using namespace facebook::velox::connector;
+
+ParquetConnector::ParquetConnector(
+    const std::string& id,
+    std::shared_ptr<const facebook::velox::config::ConfigBase> config,
+    folly::Executor* executor)
+    : Connector(id),
+      ParquetReaderConfig_(std::make_shared<ParquetReaderConfig>(config)),
+      executor_(executor) {
+  LOG(INFO) << "cudf::Parquet connector " << connectorId() << " created.";
+}
+
+std::unique_ptr<DataSource> ParquetConnector::createDataSource(
+    const std::shared_ptr<const RowType>& outputType,
+    const std::shared_ptr<ConnectorTableHandle>& tableHandle,
+    const std::unordered_map<std::string, std::shared_ptr<ColumnHandle>>&
+        columnHandles,
+    ConnectorQueryCtx* connectorQueryCtx) {
+  return std::make_unique<ParquetDataSource>(
+      outputType,
+      tableHandle,
+      columnHandles,
+      executor_,
+      connectorQueryCtx,
+      ParquetReaderConfig_);
+}
+
+std::shared_ptr<Connector> ParquetConnectorFactory::newConnector(
+    const std::string& id,
+    std::shared_ptr<const facebook::velox::config::ConfigBase> config,
+    folly::Executor* executor) {
+  return std::make_shared<ParquetConnector>(id, config, executor);
+}
+
+} // namespace facebook::velox::cudf_velox::connector::parquet

velox/experimental/cudf/connectors/parquet/ParquetConnector.h

@@ -0,0 +1,86 @@
+/*
+ * 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 "velox/connectors/Connector.h"
+#include "velox/experimental/cudf/connectors/parquet/ParquetDataSource.h"
+#include "velox/experimental/cudf/connectors/parquet/ParquetReaderConfig.h"
+#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h"
+
+#include <cudf/io/parquet.hpp>
+#include <cudf/io/types.hpp>
+#include <cudf/types.hpp>
+
+namespace facebook::velox::cudf_velox::connector::parquet {
+
+using namespace facebook::velox::connector;
+using namespace facebook::velox::config;
+
+class ParquetConnector final : public Connector {
+ public:
+  ParquetConnector(
+      const std::string& id,
+      std::shared_ptr<const ConfigBase> config,
+      folly::Executor* executor);
+
+  std::unique_ptr<DataSource> createDataSource(
+      const std::shared_ptr<const RowType>& outputType,
+      const std::shared_ptr<ConnectorTableHandle>& tableHandle,
+      const std::unordered_map<std::string, std::shared_ptr<ColumnHandle>>&
+          columnHandles,
+      ConnectorQueryCtx* connectorQueryCtx) override final;
+
+  const std::shared_ptr<const ConfigBase>& connectorConfig() const override {
+    return ParquetReaderConfig_->config();
+  }
+
+  std::unique_ptr<DataSink> createDataSink(
+      RowTypePtr /*inputType*/,
+      std::shared_ptr<
+
+          ConnectorInsertTableHandle> /*connectorInsertTableHandle*/,
+      ConnectorQueryCtx* /*connectorQueryCtx*/,
+      CommitStrategy /*commitStrategy*/) override final {
+    // TODO: Implement cudf parquet writer
+    VELOX_NYI("cudf::ParquetConnector does not yet support data sink.");
+  }
+
+  folly::Executor* executor() const override {
+    return executor_;
+  }
+
+ protected:
+  const std::shared_ptr<ParquetReaderConfig> ParquetReaderConfig_;
+  folly::Executor* executor_;
+};
+
+class ParquetConnectorFactory : public ConnectorFactory {
+ public:
+  static constexpr const char* kParquetConnectorName = "parquet";
+
+  ParquetConnectorFactory() : ConnectorFactory(kParquetConnectorName) {}
+
+  explicit ParquetConnectorFactory(const char* connectorName)
+      : ConnectorFactory(connectorName) {}
+
+  std::shared_ptr<Connector> newConnector(
+      const std::string& id,
+      std::shared_ptr<const ConfigBase> config,
+      folly::Executor* executor = nullptr) override;
+};
+
+} // namespace facebook::velox::cudf_velox::connector::parquet

velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.cpp

@@ -0,0 +1,43 @@
+/*
+ * 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 <string>
+
+#include "velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h"
+
+namespace facebook::velox::cudf_velox::connector::parquet {
+
+std::string ParquetConnectorSplit::toString() const {
+  return fmt::format("Parquet: {}", filePath);
+}
+
+std::string ParquetConnectorSplit::getFileName() const {
+  const auto i = filePath.rfind('/');
+  return i == std::string::npos ? filePath : filePath.substr(i + 1);
+}
+
+// static
+std::shared_ptr<ParquetConnectorSplit> ParquetConnectorSplit::create(
+    const folly::dynamic& obj) {
+  const auto connectorId = obj["connectorId"].asString();
+  const auto splitWeight = obj["splitWeight"].asInt();
+  const auto filePath = obj["filePath"].asString();
+
+  return std::make_shared<ParquetConnectorSplit>(
+      connectorId, filePath, splitWeight);
+}
+
+} // namespace facebook::velox::cudf_velox::connector::parquet

velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h

@@ -0,0 +1,80 @@
+/*
+ * 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 <string>
+
+#include "velox/connectors/Connector.h"
+#include "velox/dwio/common/Options.h"
+
+#include <cudf/io/types.hpp>
+
+namespace facebook::velox::cudf_velox::connector::parquet {
+
+struct ParquetConnectorSplit
+    : public facebook::velox::connector::ConnectorSplit {
+  const std::string filePath;
+  const facebook::velox::dwio::common::FileFormat fileFormat{
+      facebook::velox::dwio::common::FileFormat::PARQUET};
+  const cudf::io::source_info cudfSourceInfo;
+
+  ParquetConnectorSplit(
+      const std::string& connectorId,
+      const std::string& _filePath,
+      int64_t _splitWeight = 0)
+      : facebook::velox::connector::ConnectorSplit(connectorId, _splitWeight),
+        filePath(_filePath),
+        cudfSourceInfo({filePath}) {}
+
+  std::string toString() const override;
+  std::string getFileName() const;
+
+  const cudf::io::source_info& getCudfSourceInfo() const {
+    return cudfSourceInfo;
+  }
+
+  static std::shared_ptr<ParquetConnectorSplit> create(
+      const folly::dynamic& obj);
+};
+
+class ParquetConnectorSplitBuilder {
+ public:
+  explicit ParquetConnectorSplitBuilder(std::string filePath)
+      : filePath_{std::move(filePath)} {}
+
+  ParquetConnectorSplitBuilder& splitWeight(int64_t splitWeight) {
+    splitWeight_ = splitWeight;
+    return *this;
+  }
+
+  ParquetConnectorSplitBuilder& connectorId(const std::string& connectorId) {
+    connectorId_ = connectorId;
+    return *this;
+  }
+
+  std::shared_ptr<ParquetConnectorSplit> build() const {
+    return std::make_shared<ParquetConnectorSplit>(
+        connectorId_, filePath_, splitWeight_);
+  }
+
+ private:
+  const std::string filePath_;
+  std::string connectorId_;
+  int64_t splitWeight_{0};
+};
+
+} // namespace facebook::velox::cudf_velox::connector::parquet