chunk_writer.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 <vector>
 
#include "graphar/fwd.h"
#include "graphar/writer_util.h"
 
// forward declaration
namespace arrow {
class Table;
}
 
namespace graphar {
 
class VertexPropertyWriter {
 public:
  explicit VertexPropertyWriter(
      const std::shared_ptr<VertexInfo>& vertex_info, const std::string& prefix,
      const std::shared_ptr<WriterOptions>& options =
          WriterOptions::DefaultWriterOption(),
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  inline void SetValidateLevel(const ValidateLevel& validate_level) {
    if (validate_level == ValidateLevel::default_validate) {
      return;
    }
    validate_level_ = validate_level;
  }
 
  inline ValidateLevel GetValidateLevel() const { return validate_level_; }
 
  Status WriteVerticesNum(
      const IdType& count,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteChunk(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::shared_ptr<PropertyGroup>& property_group, IdType chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteLabelChunk(
      const std::shared_ptr<arrow::Table>& input_table, IdType chunk_index,
      FileType file_type,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteChunk(
      const std::shared_ptr<arrow::Table>& input_table, IdType chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteTable(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::shared_ptr<PropertyGroup>& property_group,
      IdType start_chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteTable(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType start_chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteLabelTable(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType start_chunk_index, FileType file_type,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Result<std::shared_ptr<arrow::Table>> GetLabelTable(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::vector<std::string>& labels) const;
 
  Result<std::shared_ptr<arrow::Table>> GetLabelTableAndRandomlyAddLabels(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::vector<std::string>& labels) const;
 
  static Result<std::shared_ptr<VertexPropertyWriter>> Make(
      const std::shared_ptr<VertexInfo>& vertex_info, const std::string& prefix,
      const std::shared_ptr<WriterOptions>& options,
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  static Result<std::shared_ptr<VertexPropertyWriter>> Make(
      const std::shared_ptr<VertexInfo>& vertex_info, const std::string& prefix,
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  static Result<std::shared_ptr<VertexPropertyWriter>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& type,
      const std::shared_ptr<WriterOptions>& options,
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  static Result<std::shared_ptr<VertexPropertyWriter>> Make(
      const std::shared_ptr<GraphInfo>& graph_info, const std::string& type,
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  void setWriterOptions(const std::shared_ptr<WriterOptions>& options) {
    options_ = options;
  }
 
  Result<std::shared_ptr<arrow::Table>> AddIndexColumn(
      const std::shared_ptr<arrow::Table>& table, IdType chunk_index,
      IdType chunk_size) const;
 
 private:
  Status validate(const IdType& count, ValidateLevel validate_level) const;
 
  Status validate(const std::shared_ptr<PropertyGroup>& property_group,
                  IdType chunk_index, ValidateLevel validate_level) const;
 
  Status validate(const std::shared_ptr<arrow::Table>& input_table,
                  const std::shared_ptr<PropertyGroup>& property_group,
                  IdType chunk_index, ValidateLevel validate_level) const;
 
 private:
  std::shared_ptr<VertexInfo> vertex_info_;
  std::string prefix_;
  std::shared_ptr<FileSystem> fs_;
  ValidateLevel validate_level_;
  std::shared_ptr<WriterOptions> options_;
};
 
class EdgeChunkWriter {
 public:
  explicit EdgeChunkWriter(
      const std::shared_ptr<EdgeInfo>& edge_info, const std::string& prefix,
      AdjListType adj_list_type,
      const std::shared_ptr<WriterOptions>& options =
          WriterOptions::DefaultWriterOption(),
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  void SetValidateLevel(const ValidateLevel& validate_level) {
    if (validate_level == ValidateLevel::default_validate) {
      return;
    }
    validate_level_ = validate_level;
  }
 
  inline ValidateLevel GetValidateLevel() const { return validate_level_; }
 
  Status WriteEdgesNum(
      IdType vertex_chunk_index, const IdType& count,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteVerticesNum(
      const IdType& count,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteOffsetChunk(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteAdjListChunk(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WritePropertyChunk(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::shared_ptr<PropertyGroup>& property_group,
      IdType vertex_chunk_index, IdType chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WritePropertyChunk(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteChunk(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType chunk_index,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteAdjListTable(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType start_chunk_index = 0,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WritePropertyTable(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::shared_ptr<PropertyGroup>& property_group,
      IdType vertex_chunk_index, IdType start_chunk_index = 0,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WritePropertyTable(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType start_chunk_index = 0,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status WriteTable(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType start_chunk_index = 0,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status SortAndWriteAdjListTable(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType start_chunk_index = 0,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status SortAndWritePropertyTable(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::shared_ptr<PropertyGroup>& property_group,
      IdType vertex_chunk_index, IdType start_chunk_index = 0,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status SortAndWritePropertyTable(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType start_chunk_index = 0,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  Status SortAndWriteTable(
      const std::shared_ptr<arrow::Table>& input_table,
      IdType vertex_chunk_index, IdType start_chunk_index = 0,
      ValidateLevel validate_level = ValidateLevel::default_validate) const;
 
  static Result<std::shared_ptr<EdgeChunkWriter>> Make(
      const std::shared_ptr<EdgeInfo>& edge_info, const std::string& prefix,
      AdjListType adj_list_type, const std::shared_ptr<WriterOptions>& options,
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  static Result<std::shared_ptr<EdgeChunkWriter>> Make(
      const std::shared_ptr<EdgeInfo>& edge_info, const std::string& prefix,
      AdjListType adj_list_type,
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  static Result<std::shared_ptr<EdgeChunkWriter>> 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, const std::shared_ptr<WriterOptions>& options,
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
  static Result<std::shared_ptr<EdgeChunkWriter>> 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,
      const ValidateLevel& validate_level = ValidateLevel::no_validate);
 
 private:
  Status validate(IdType count_or_index1, IdType count_or_index2,
                  ValidateLevel validate_level) const;
 
  Status validate(const std::shared_ptr<PropertyGroup>& property_group,
                  IdType vertex_chunk_index, IdType chunk_index,
                  ValidateLevel validate_level) const;
 
  Status validate(const std::shared_ptr<arrow::Table>& input_table,
                  IdType vertex_chunk_index,
                  ValidateLevel validate_level) const;
 
  Status validate(const std::shared_ptr<arrow::Table>& input_table,
                  IdType vertex_chunk_index, IdType chunk_index,
                  ValidateLevel validate_level) const;
 
  Status validate(const std::shared_ptr<arrow::Table>& input_table,
                  const std::shared_ptr<PropertyGroup>& property_group,
                  IdType vertex_chunk_index, IdType chunk_index,
                  ValidateLevel validate_level) const;
 
  Result<std::shared_ptr<arrow::Table>> getOffsetTable(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::string& column_name, IdType vertex_chunk_index) const;
 
  static std::string getSortColumnName(AdjListType adj_list_type);
 
  static Result<std::shared_ptr<arrow::Table>> sortTable(
      const std::shared_ptr<arrow::Table>& input_table,
      const std::string& column_name);
 
 private:
  std::shared_ptr<EdgeInfo> edge_info_;
  IdType vertex_chunk_size_;
  IdType chunk_size_;
  AdjListType adj_list_type_;
  std::string prefix_;
  std::shared_ptr<FileSystem> fs_;
  ValidateLevel validate_level_;
  std::shared_ptr<WriterOptions> options_;
};
 
}  // namespace graphar