chunk_reader.h

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 <memory>
#include <string>
#include <utility>
#include <vector>
 
#include "graphar/fwd.h"
#include "graphar/reader_util.h"
#include "graphar/status.h"
 
// forward declaration
namespace arrow {
class Array;
class Schema;
class Table;
}  // namespace arrow
 
namespace graphar {
 
class VertexPropertyArrowChunkReader {
 public:
  VertexPropertyArrowChunkReader(
      const std::shared_ptr<VertexInfo>& vertex_info,
      const std::shared_ptr<PropertyGroup>& property_group,
      const std::string& prefix, const util::FilterOptions& options = {});
  VertexPropertyArrowChunkReader(
      const std::shared_ptr<VertexInfo>& vertex_info,
      const std::shared_ptr<PropertyGroup>& property_group,
      const std::vector<std::string>& property_names, const std::string& prefix,
      const util::FilterOptions& options = {});
 
  VertexPropertyArrowChunkReader() : vertex_info_(nullptr), prefix_("") {}
 
  VertexPropertyArrowChunkReader(const std::shared_ptr<VertexInfo>& vertex_info,
                                 const std::vector<std::string>& labels,
                                 const std::string& prefix,
                                 const util::FilterOptions& options = {});
  Status seek(IdType id);
 
  Result<std::shared_ptr<arrow::Table>> GetChunk(
      GetChunkVersion version = GetChunkVersion::AUTO);
  Result<std::shared_ptr<arrow::Table>> GetLabelChunk();
  Status next_chunk();
 
  IdType GetChunkNum() const noexcept { return chunk_num_; }
 
  void Filter(util::Filter filter = nullptr);
 
  void Select(util::ColumnNames column_names = std::nullopt);
 
  static Result<std::shared_ptr<VertexPropertyArrowChunkReader>> Make(
      const std::shared_ptr<VertexInfo>& vertex_info,
      const std::shared_ptr<PropertyGroup>& property_group,
      const std::string& prefix, const util::FilterOptions& options = {});
 
  static Result<std::shared_ptr<VertexPropertyArrowChunkReader>> Make(
      const std::shared_ptr<VertexInfo>& vertex_info,
      const std::shared_ptr<PropertyGroup>& property_group,
      const std::vector<std::string>& property_names, const std::string& prefix,
      const util::FilterOptions& options = {});
 
  static Result<std::shared_ptr<VertexPropertyArrowChunkReader>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& type,
      const std::shared_ptr<PropertyGroup>& property_group,
      const util::FilterOptions& options = {});
 
  static Result<std::shared_ptr<VertexPropertyArrowChunkReader>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& type,
      const std::string& property_name,
      const util::FilterOptions& options = {});
  static Result<std::shared_ptr<VertexPropertyArrowChunkReader>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& type,
      const std::vector<std::string>& property_names_or_labels,
      const SelectType select_type, const util::FilterOptions& options = {});
 
  static Result<std::shared_ptr<VertexPropertyArrowChunkReader>> Make(
      const std::shared_ptr<VertexInfo>& vertex_info,
      const std::vector<std::string>& labels, const std::string& prefix,
      const util::FilterOptions& options = {});
 
  static Result<std::shared_ptr<VertexPropertyArrowChunkReader>>
  MakeForProperties(const std::shared_ptr<GraphInfo>& graph_info,
                    const std::string& type,
                    const std::vector<std::string>& property_names,
                    const util::FilterOptions& options = {});
 
  static Result<std::shared_ptr<VertexPropertyArrowChunkReader>> MakeForLabels(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& type,
      const std::vector<std::string>& labels,
      const util::FilterOptions& options = {});
 
 private:
  Result<std::shared_ptr<arrow::Table>> GetChunkV1();
  Result<std::shared_ptr<arrow::Table>> GetChunkV2();
 
 private:
  std::shared_ptr<VertexInfo> vertex_info_;
  std::shared_ptr<PropertyGroup> property_group_;
  std::vector<std::string> property_names_;
  std::string prefix_;
  std::vector<std::string> labels_;
  IdType chunk_index_;
  IdType seek_id_;
  IdType chunk_num_;
  IdType vertex_num_;
  std::shared_ptr<arrow::Schema> schema_;
  std::shared_ptr<arrow::Table> chunk_table_;
  util::FilterOptions filter_options_;
  std::shared_ptr<FileSystem> fs_;
};
 
class AdjListArrowChunkReader {
 public:
  using range_t = std::pair<IdType, IdType>;
  AdjListArrowChunkReader(const std::shared_ptr<EdgeInfo>& edge_info,
                          AdjListType adj_list_type, const std::string& prefix);
 
  AdjListArrowChunkReader(const AdjListArrowChunkReader& other);
 
  Status seek_src(IdType id);
 
  Status seek_dst(IdType offset);
 
  Status seek(IdType offset);
 
  Result<std::shared_ptr<arrow::Table>> GetChunk();
 
  Result<IdType> GetRowNumOfChunk();
 
  Status next_chunk();
 
  Status seek_chunk_index(IdType vertex_chunk_index, IdType chunk_index = 0);
 
  static Result<std::shared_ptr<AdjListArrowChunkReader>> Make(
      const std::shared_ptr<EdgeInfo>& edge_info, AdjListType adj_list_type,
      const std::string& prefix);
 
  static Result<std::shared_ptr<AdjListArrowChunkReader>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& src_type,
      const std::string& edge_type, const std::string& dst_type,
      AdjListType adj_list_type);
 
 private:
  Status initOrUpdateEdgeChunkNum();
 
 private:
  std::shared_ptr<EdgeInfo> edge_info_;
  AdjListType adj_list_type_;
  std::string prefix_;
  IdType vertex_chunk_index_, chunk_index_;
  IdType seek_offset_;
  std::shared_ptr<arrow::Table> chunk_table_;
  IdType vertex_chunk_num_, chunk_num_;
  std::string base_dir_;
  std::shared_ptr<FileSystem> fs_;
};
 
class AdjListOffsetArrowChunkReader {
 public:
  using range_t = std::pair<IdType, IdType>;
  AdjListOffsetArrowChunkReader(const std::shared_ptr<EdgeInfo>& edge_info,
                                AdjListType adj_list_type,
                                const std::string& prefix);
 
  Status seek(IdType id);
 
  Result<std::shared_ptr<arrow::Array>> GetChunk();
 
  Status next_chunk();
 
  IdType GetChunkIndex() const noexcept { return chunk_index_; }
 
  static Result<std::shared_ptr<AdjListOffsetArrowChunkReader>> Make(
      const std::shared_ptr<EdgeInfo>& edge_info, AdjListType adj_list_type,
      const std::string& prefix);
 
  static Result<std::shared_ptr<AdjListOffsetArrowChunkReader>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& src_type,
      const std::string& edge_type, const std::string& dst_type,
      AdjListType adj_list_type);
 
 private:
  std::shared_ptr<EdgeInfo> edge_info_;
  AdjListType adj_list_type_;
  std::string prefix_;
  IdType chunk_index_;
  IdType seek_id_;
  std::shared_ptr<arrow::Table> chunk_table_;
  IdType vertex_chunk_num_;
  IdType vertex_chunk_size_;
  std::string base_dir_;
  std::shared_ptr<FileSystem> fs_;
};
 
class AdjListPropertyArrowChunkReader {
 public:
  using range_t = std::pair<IdType, IdType>;
  AdjListPropertyArrowChunkReader(
      const std::shared_ptr<EdgeInfo>& edge_info,
      const std::shared_ptr<PropertyGroup>& property_group,
      AdjListType adj_list_type, const std::string prefix,
      const util::FilterOptions& options = {});
 
  AdjListPropertyArrowChunkReader(const AdjListPropertyArrowChunkReader& other);
 
  Status seek_src(IdType id);
 
  Status seek_dst(IdType id);
 
  Status seek(IdType offset);
 
  Result<std::shared_ptr<arrow::Table>> GetChunk();
 
  Status next_chunk();
 
  Status seek_chunk_index(IdType vertex_chunk_index, IdType chunk_index = 0);
 
  void Filter(util::Filter filter = nullptr);
 
  void Select(util::ColumnNames column_names = std::nullopt);
 
  static Result<std::shared_ptr<AdjListPropertyArrowChunkReader>> Make(
      const std::shared_ptr<EdgeInfo>& edge_info,
      const std::shared_ptr<PropertyGroup>& property_group,
      AdjListType adj_list_type, const std::string& prefix,
      const util::FilterOptions& options = {});
 
  static Result<std::shared_ptr<AdjListPropertyArrowChunkReader>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& src_type,
      const std::string& edge_type, const std::string& dst_type,
      const std::shared_ptr<PropertyGroup>& property_group,
      AdjListType adj_list_type, const util::FilterOptions& options = {});
 
  static Result<std::shared_ptr<AdjListPropertyArrowChunkReader>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& src_type,
      const std::string& edge_type, const std::string& dst_type,
      const std::string& property_name, AdjListType adj_list_type,
      const util::FilterOptions& options = {});
 
 private:
  Status initOrUpdateEdgeChunkNum();
 
 private:
  std::shared_ptr<EdgeInfo> edge_info_;
  std::shared_ptr<PropertyGroup> property_group_;
  AdjListType adj_list_type_;
  std::string prefix_;
  IdType vertex_chunk_index_, chunk_index_;
  IdType seek_offset_;
  std::shared_ptr<arrow::Schema> schema_;
  std::shared_ptr<arrow::Table> chunk_table_;
  util::FilterOptions filter_options_;
  IdType vertex_chunk_num_, chunk_num_;
  std::string base_dir_;
  std::shared_ptr<FileSystem> fs_;
};
}  // namespace graphar