[Feat][Format] Add internal id column to vertex payload file (#264)

---------

Signed-off-by: acezen <qiaozi.zwb@alibaba-inc.com>

Author: acezen

.github/workflows/ci.yml

@@ -36,10 +36,10 @@ jobs:
             -P /tmp/
         sudo apt-get install -y /tmp/apache-arrow-apt-source-latest-"$(lsb_release --codename --short)".deb
         sudo apt-get update -y
-        sudo apt install -y libarrow-dev=14.0.0-1 \
-                            libarrow-dataset-dev=14.0.0-1 \
-                            libarrow-acero-dev=14.0.0-1 \
-                            libparquet-dev=14.0.0-1
+        sudo apt install -y libarrow-dev=14.0.1-1 \
+                            libarrow-dataset-dev=14.0.1-1 \
+                            libarrow-acero-dev=14.0.1-1 \
+                            libparquet-dev=14.0.1-1
         sudo apt-get install -y libboost-graph-dev ccache libcurl4-openssl-dev
 
     - name: CMake

README.rst

@@ -98,6 +98,8 @@ Take the "person" vertex table as an example, if the chunk size is set to be 500
   :align: center
   :alt: vertex physical table
 
+**Note**: For efficiently utilize the filter push-down of the payload file format like Parquet, the internal vertex id is stored in the payload file as a column. And since the internal vertex id is continuous, the payload file format can use the delta encoding for the internal vertex id column, which would not bring too much overhead for the storage.
+
 Edges in GraphAr
 ^^^^^^^^^^^^^^^^
 

cpp/examples/mid_level_reader_example.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include "./config.h"
 #include "gar/reader/arrow_chunk_reader.h"
 #include "gar/util/expression.h"
+#include "gar/util/general_params.h"
 
 void vertex_property_chunk_reader(const GAR_NAMESPACE::GraphInfo& graph_info) {
   // constuct reader
@@ -36,29 +37,37 @@ void vertex_property_chunk_reader(const GAR_NAMESPACE::GraphInfo& graph_info) {
   // use reader
   auto result = reader.GetChunk();
   ASSERT(!result.has_error());
-  auto range = reader.GetRange().value();
   std::cout << "chunk number: " << reader.GetChunkNum() << std::endl;
-  std::cout << "range of fisrt vertex property chunk: " << range.first << " "
-            << range.second << std::endl;
   auto table = result.value();
   std::cout << "rows number of first vertex property chunk: "
             << table->num_rows() << std::endl;
   std::cout << "schema of first vertex property chunk: " << std::endl
             << table->schema()->ToString() << std::endl;
+  auto index_col =
+      table->GetColumnByName(GAR_NAMESPACE::GeneralParams::kVertexIndexCol);
+  ASSERT(index_col != nullptr);
+  std::cout << "Internal id column: " << index_col->ToString() << " "
+            << std::endl;
   // seek vertex id
   ASSERT(reader.seek(100).ok());
   result = reader.GetChunk();
   ASSERT(!result.has_error());
-  range = reader.GetRange().value();
-  std::cout << "range of vertex property chunk for vertex id 100: "
-            << range.first << " " << range.second << std::endl;
+  table = result.value();
+  index_col =
+      table->GetColumnByName(GAR_NAMESPACE::GeneralParams::kVertexIndexCol);
+  ASSERT(index_col != nullptr);
+  std::cout << "Internal id column of vertex property chunk for vertex id 100: "
+            << index_col->ToString() << " " << std::endl;
   // next chunk
   ASSERT(reader.next_chunk().ok());
   result = reader.GetChunk();
   ASSERT(!result.has_error());
-  range = reader.GetRange().value();
-  std::cout << "range of next vertex property chunk: " << range.first << " "
-            << range.second << std::endl;
+  table = result.value();
+  index_col =
+      table->GetColumnByName(GAR_NAMESPACE::GeneralParams::kVertexIndexCol);
+  ASSERT(index_col != nullptr);
+  std::cout << "Internal id column of next chunk: " << index_col->ToString()
+            << " " << std::endl;
 
   // reader with filter pushdown
   auto filter = GAR_NAMESPACE::_Equal(GAR_NAMESPACE::_Property("gender"),

cpp/include/gar/reader/arrow_chunk_reader.h

@@ -100,13 +100,6 @@ class VertexPropertyArrowChunkReader {
    */
   Result<std::shared_ptr<arrow::Table>> GetChunk() noexcept;
 
-  /**
-   * @brief Get the vertex id range of current chunk.
-   *
-   * @return Result: std::pair<begin_id, end_id> or error.
-   */
-  Result<std::pair<IdType, IdType>> GetRange() noexcept;
-
   /**
    * @brief Sets chunk position indicator to next chunk.
    *

cpp/include/gar/writer/arrow_chunk_writer.h

@@ -227,6 +227,10 @@ class VertexPropertyWriter {
                   ValidateLevel validate_level) const noexcept;
 
  private:
+  Result<std::shared_ptr<arrow::Table>> addIndexColumn(
+      const std::shared_ptr<arrow::Table>& table, IdType chunk_index,
+      IdType chunk_size) const noexcept;
+
   VertexInfo vertex_info_;
   std::string prefix_;
   std::shared_ptr<FileSystem> fs_;

cpp/src/arrow_chunk_reader.cc

@@ -36,22 +36,6 @@ VertexPropertyArrowChunkReader::GetChunk() noexcept {
   return chunk_table_->Slice(row_offset);
 }
 
-Result<std::pair<IdType, IdType>>
-VertexPropertyArrowChunkReader::GetRange() noexcept {
-  if (chunk_table_ == nullptr) {
-    return Status::Invalid(
-        "The chunk table is not initialized, please call "
-        "GetChunk() first.");
-  }
-  const auto chunk_size = vertex_info_.GetChunkSize();
-  IdType row_offset = seek_id_ - chunk_index_ * chunk_size;
-  bool is_last_chunk = (chunk_index_ == chunk_num_ - 1);
-  const auto curr_chunk_size =
-      is_last_chunk ? (vertex_num_ - chunk_index_ * chunk_size) : chunk_size;
-
-  return std::make_pair(seek_id_, seek_id_ + curr_chunk_size - row_offset);
-}
-
 void VertexPropertyArrowChunkReader::Filter(util::Filter filter) {
   filter_options_.filter = filter;
 }

cpp/src/arrow_chunk_writer.cc

@@ -198,10 +198,15 @@ Status VertexPropertyWriter::WriteChunk(
   GAR_RETURN_NOT_OK(
       validate(input_table, property_group, chunk_index, validate_level));
   auto file_type = property_group.GetFileType();
-
-  std::vector<int> indices;
-  indices.clear();
   auto schema = input_table->schema();
+  int indice = schema->GetFieldIndex(GeneralParams::kVertexIndexCol);
+  if (indice == -1) {
+    return Status::Invalid("The internal id Column named ",
+                           GeneralParams::kVertexIndexCol,
+                           " does not exist in the input table.");
+  }
+
+  std::vector<int> indices({indice});
   for (auto& property : property_group.GetProperties()) {
     int indice = schema->GetFieldIndex(property.name);
     if (indice == -1) {
@@ -251,13 +256,37 @@ Status VertexPropertyWriter::WriteTable(
     const std::shared_ptr<arrow::Table>& input_table, IdType start_chunk_index,
     ValidateLevel validate_level) const noexcept {
   auto property_groups = vertex_info_.GetPropertyGroups();
+  GAR_ASSIGN_OR_RAISE(auto table_with_index,
+                      addIndexColumn(input_table, start_chunk_index,
+                                     vertex_info_.GetChunkSize()));
   for (auto& property_group : property_groups) {
-    GAR_RETURN_NOT_OK(WriteTable(input_table, property_group, start_chunk_index,
-                                 validate_level));
+    GAR_RETURN_NOT_OK(WriteTable(table_with_index, property_group,
+                                 start_chunk_index, validate_level));
   }
   return Status::OK();
 }
 
+Result<std::shared_ptr<arrow::Table>> VertexPropertyWriter::addIndexColumn(
+    const std::shared_ptr<arrow::Table>& table, IdType chunk_index,
+    IdType chunk_size) const noexcept {
+  arrow::Int64Builder array_builder;
+  RETURN_NOT_ARROW_OK(array_builder.Reserve(chunk_size));
+  int64_t length = table->num_rows();
+  for (IdType i = 0; i < length; i++) {
+    RETURN_NOT_ARROW_OK(array_builder.Append(chunk_index * chunk_size + i));
+  }
+  std::shared_ptr<arrow::Array> array;
+  RETURN_NOT_ARROW_OK(array_builder.Finish(&array));
+  std::shared_ptr<arrow::ChunkedArray> chunked_array =
+      std::make_shared<arrow::ChunkedArray>(array);
+  GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(
+      auto ret, table->AddColumn(0,
+                                 arrow::field(GeneralParams::kVertexIndexCol,
+                                              arrow::int64(), false),
+                                 chunked_array));
+  return ret;
+}
+
 // implementations for EdgeChunkWriter
 
 // Check if the operation of writing number or copying a file is allowed.

cpp/test/test_arrow_chunk_reader.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include "./util.h"
 #include "gar/reader/arrow_chunk_reader.h"
 #include "gar/util/expression.h"
+#include "gar/util/general_params.h"
 
 #define CATCH_CONFIG_MAIN
 #include <catch2/catch.hpp>
@@ -53,37 +54,33 @@ TEST_CASE("test_vertex_property_arrow_chunk_reader") {
   auto reader = maybe_reader.value();
   auto result = reader.GetChunk();
   REQUIRE(!result.has_error());
-  auto range = reader.GetRange().value();
   auto table = result.value();
   REQUIRE(table->num_rows() == 100);
-  REQUIRE(range.first == 0);
-  REQUIRE(range.second == 100);
+  REQUIRE(table->GetColumnByName(
+              GAR_NAMESPACE::GeneralParams::kVertexIndexCol) != nullptr);
 
   // seek
   REQUIRE(reader.seek(100).ok());
   result = reader.GetChunk();
   REQUIRE(!result.has_error());
-  range = reader.GetRange().value();
   table = result.value();
   REQUIRE(table->num_rows() == 100);
-  REQUIRE(range.first == 100);
-  REQUIRE(range.second == 200);
+  REQUIRE(table->GetColumnByName(
+              GAR_NAMESPACE::GeneralParams::kVertexIndexCol) != nullptr);
   REQUIRE(reader.next_chunk().ok());
   result = reader.GetChunk();
   REQUIRE(!result.has_error());
-  range = reader.GetRange().value();
   table = result.value();
   REQUIRE(table->num_rows() == 100);
-  REQUIRE(range.first == 200);
-  REQUIRE(range.second == 300);
+  REQUIRE(table->GetColumnByName(
+              GAR_NAMESPACE::GeneralParams::kVertexIndexCol) != nullptr);
   REQUIRE(reader.seek(900).ok());
   result = reader.GetChunk();
   REQUIRE(!result.has_error());
-  range = reader.GetRange().value();
   table = result.value();
   REQUIRE(table->num_rows() == 3);
-  REQUIRE(range.first == 900);
-  REQUIRE(range.second == 903);
+  REQUIRE(table->GetColumnByName(
+              GAR_NAMESPACE::GeneralParams::kVertexIndexCol) != nullptr);
   REQUIRE(reader.GetChunkNum() == 10);
   REQUIRE(reader.next_chunk().IsIndexError());
 
@@ -119,10 +116,7 @@ TEST_CASE("test_vertex_property_pushdown") {
       auto result = reader.GetChunk();
       REQUIRE(!result.has_error());
       table = result.value();
-      auto [start, end] = reader.GetRange().value();
-      std::cout << "Chunk: " << idx << ",\tNums: " << table->num_rows() << "/"
-                << chunk_size << ",\tRange: (" << start << ", " << end << "]"
-                << '\n';
+      std::cout << "Chunk: " << idx << ",\tNums: " << table->num_rows() << '\n';
       idx++;
       sum += table->num_rows();
     } while (!reader.next_chunk().IsIndexError());

docs/images/vertex_physical_table.png

@@ -57,6 +57,8 @@ Take the "person" vertex table as an example, if the chunk size is set to be 500
    :alt: vertex physical table
 
 
+**Note**: For efficiently utilize the filter push-down of the payload file format like Parquet, the internal vertex id is stored in the payload file as a column. And since the internal vertex id is continuous, the payload file format can use the delta encoding for the internal vertex id column, which would not bring too much overhead for the storage.
+
 Edges in GraphAr
 ------------------------