odps_fragment_loader.h

Go to the documentation of this file.

 
#ifndef STORAGES_RT_MUTABLE_GRAPH_LOADER_ODPS_FRAGMENT_LOADER_H_
#define STORAGES_RT_MUTABLE_GRAPH_LOADER_ODPS_FRAGMENT_LOADER_H_
 
#include <arrow/api.h>
#include <arrow/csv/api.h>
#include <arrow/io/api.h>
#include <arrow/util/uri.h>
 
#include <boost/convert.hpp>
#include <boost/convert/strtol.hpp>
#include <charconv>
 
#include "arrow/util/value_parsing.h"
#include "common/configuration.h"
#include "flex/storages/rt_mutable_graph/loader/abstract_arrow_fragment_loader.h"
#include "flex/storages/rt_mutable_graph/loader/basic_fragment_loader.h"
#include "flex/storages/rt_mutable_graph/loader/i_fragment_loader.h"
#include "flex/storages/rt_mutable_graph/loader/loader_factory.h"
#include "flex/storages/rt_mutable_graph/loading_config.h"
#include "flex/storages/rt_mutable_graph/mutable_property_fragment.h"
#include "flex/third_party/httplib.h"
#include "grape/util.h"
#include "storage_api.hpp"
#include "storage_api_arrow.hpp"
 
using apsara::odps::sdk::AliyunAccount;
using apsara::odps::sdk::Configuration;
using apsara::odps::sdk::storage_api::ReadRowsReq;
using apsara::odps::sdk::storage_api::SessionReq;
using apsara::odps::sdk::storage_api::SessionStatus;
using apsara::odps::sdk::storage_api::SplitOptions;
using apsara::odps::sdk::storage_api::TableBatchScanReq;
using apsara::odps::sdk::storage_api::TableBatchScanResp;
using apsara::odps::sdk::storage_api::TableBatchWriteReq;
using apsara::odps::sdk::storage_api::TableBatchWriteResp;
using apsara::odps::sdk::storage_api::TableIdentifier;
using apsara::odps::sdk::storage_api::WriteRowsReq;
using apsara::odps::sdk::storage_api::arrow_adapter::ArrowClient;
using apsara::odps::sdk::storage_api::arrow_adapter::Reader;
 
namespace gs {
 
class ODPSReadClient {
 public:
  static constexpr const int CONNECTION_TIMEOUT = 5;
  static constexpr const int READ_WRITE_TIMEOUT = 10;
  static constexpr const size_t MAX_RETRY = 10;
  ODPSReadClient();
 
  ~ODPSReadClient();
 
  void init();
 
  void CreateReadSession(std::string* session_id, int* split_count,
                         const TableIdentifier& table_identifier,
                         const std::vector<std::string>& selected_cols,
                         const std::vector<std::string>& partition_cols,
                         const std::vector<std::string>& selected_partitions);
 
  std::shared_ptr<arrow::Table> ReadTable(const std::string& session_id,
                                          int split_count,
                                          const TableIdentifier& table_id,
                                          int thread_num) const;
 
  std::shared_ptr<ArrowClient> GetArrowClient() const;
 
 private:
  TableBatchScanResp createReadSession(
      const TableIdentifier& table_identifier,
      const std::vector<std::string>& selected_cols,
      const std::vector<std::string>& partition_cols,
      const std::vector<std::string>& selected_partitions);
 
  TableBatchScanResp getReadSession(std::string session_id,
                                    const TableIdentifier& table_identifier);
 
  void getReadSessionStatus(const std::string& session_id, int* split_count,
                            const TableIdentifier& table_identifier);
 
  void producerRoutine(
      const std::string& session_id, const TableIdentifier& table_identifier,
      std::vector<std::vector<std::shared_ptr<arrow::RecordBatch>>>&
          all_batches_,
      std::vector<int>&& indices) const;
 
  bool readRows(std::string session_id, const TableIdentifier& table_identifier,
                std::vector<std::shared_ptr<arrow::RecordBatch>>& res_batches,
                int split_index) const;
 
 private:
  // odps table related
  std::string access_id_;
  std::string access_key_;
  std::string odps_endpoint_;
  std::string tunnel_endpoint_;
  std::string output_directory_;
  std::shared_ptr<ArrowClient> arrow_client_ptr_;
  size_t MAX_PRODUCER_NUM = 8;
};
 
class ODPSStreamRecordBatchSupplier : public IRecordBatchSupplier {
 public:
  ODPSStreamRecordBatchSupplier(label_t label_id, const std::string& file_path,
                                const ODPSReadClient& odps_table_reader,
                                const std::string& session_id, int split_count,
                                TableIdentifier table_identifier, int worker_id,
                                int worker_num);
 
  std::shared_ptr<arrow::RecordBatch> GetNextBatch() override;
 
 private:
  std::string file_path_;
  const ODPSReadClient& odps_read_client_;
  std::string session_id_;
  int split_count_;
 
  TableIdentifier table_identifier_;
 
  int32_t cur_split_index_;
  int worker_num_;
  ReadRowsReq read_rows_req_;
  std::shared_ptr<Reader> cur_batch_reader_;
};
 
class ODPSTableRecordBatchSupplier : public IRecordBatchSupplier {
 public:
  ODPSTableRecordBatchSupplier(label_t label_id, const std::string& file_path,
                               const ODPSReadClient& odps_table_reader,
                               const std::string& session_id, int split_count,
                               TableIdentifier table_identifier,
                               int thread_num);
 
  std::shared_ptr<arrow::RecordBatch> GetNextBatch() override;
 
 private:
  std::string file_path_;
  const ODPSReadClient& odps_read_client_;
  std::string session_id_;
  TableIdentifier table_identifier_;
 
  std::shared_ptr<arrow::Table> table_;
  std::shared_ptr<arrow::TableBatchReader> reader_;
};
 
/*
 * ODPSFragmentLoader is used to load graph data from ODPS Table.
 * It fetch the data via ODPS tunnel/halo API.
 * You need to set the following environment variables:
 * 1. ODPS_ACCESS_ID
 * 2. ODPS_ACCESS_KEY
 * 3. ODPS_ENDPOINT
 * 4. ODPS_TUNNEL_ENDPOINT(optional)
 */
class ODPSFragmentLoader : public AbstractArrowFragmentLoader {
 public:
  ODPSFragmentLoader(const std::string& work_dir, const Schema& schema,
                     const LoadingConfig& loading_config)
      : AbstractArrowFragmentLoader(work_dir, schema, loading_config) {}
 
  static std::shared_ptr<IFragmentLoader> Make(
      const std::string& work_dir, const Schema& schema,
      const LoadingConfig& loading_config);
 
  ~ODPSFragmentLoader() {}
 
  Result<bool> LoadFragment() override;
 
 private:
  void init();
 
  void parseLocation(const std::string& odps_table_path,
                     TableIdentifier& table_identifier,
                     std::vector<std::string>& partition_names,
                     std::vector<std::string>& selected_partitions);
 
  void loadVertices();
 
  void loadEdges();
 
  void addVertices(label_t v_label_id, const std::vector<std::string>& v_files);
 
  void addEdges(label_t src_label_id, label_t dst_label_id, label_t e_label_id,
                const std::vector<std::string>& e_files);
 
  std::vector<std::string> columnMappingsToSelectedCols(
      const std::vector<std::tuple<size_t, std::string, std::string>>&
          column_mappings);
 
  ODPSReadClient odps_read_client_;
 
  static const bool registered_;
};
 
}  // namespace gs
 
#endif  // STORAGES_RT_MUTABLE_GRAPH_LOADER_ODPS_FRAGMENT_LOADER_H_