Fixes typecheck and dispatch in gar fragment loader

Signed-off-by: Tao He <[email protected]>
v6d-io · Jan 20, 2024 · bdf1914 · bdf1914
1 parent ceab36b
commit bdf1914
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 7 deletions.
diff --git a/modules/graph/loader/gar_fragment_loader_impl.h b/modules/graph/loader/gar_fragment_loader_impl.h
@@ -135,14 +135,14 @@ GARFragmentLoader<OID_T, VID_T, VERTEX_MAP_T>::LoadEdgeTables() {
     const auto& edge_info = edge.second;
     if (edge_info.ContainAdjList(
             GraphArchive::AdjListType::ordered_by_source)) {
-      loadEdgeTableOfLabel(edge_info,
-                           GraphArchive::AdjListType::ordered_by_source);
+      BOOST_LEAF_CHECK(loadEdgeTableOfLabel(
+          edge_info, GraphArchive::AdjListType::ordered_by_source));
     }
     if (this->directed_) {
       if (edge_info.ContainAdjList(
               GraphArchive::AdjListType::ordered_by_dest)) {
-        loadEdgeTableOfLabel(edge_info,
-                             GraphArchive::AdjListType::ordered_by_dest);
+        BOOST_LEAF_CHECK(loadEdgeTableOfLabel(
+            edge_info, GraphArchive::AdjListType::ordered_by_dest));
       }
     }
   }

diff --git a/modules/graph/test/arrow_fragment_extend_test.cc b/modules/graph/test/arrow_fragment_extend_test.cc
@@ -129,20 +129,20 @@ bool Validate(vineyard::Client& client, const grape::CommSpec& comm_spec,
       LOG(ERROR) << "different inner edges number";
       return false;
     }
-    for (int j = 0; j < true_ivs[i].size(); ++j) {
+    for (size_t j = 0; j < true_ivs[i].size(); ++j) {
       if (true_ivs[i][j] != test_ivs[i][j]) {
         LOG(ERROR) << "ground-truth v is " << true_ivs[i][j]
                    << "and program is " << test_ivs[i][j];
         return false;
       }
     }
-    for (int j = 0; j < true_ies[i].size(); ++j) {
+    for (size_t j = 0; j < true_ies[i].size(); ++j) {
       if (true_ies[i][j] != test_ies[i][j]) {
         LOG(ERROR) << "different inner edge";
         return false;
       }
     }
-    for (int j = 0; j < true_oes[i].size(); ++j) {
+    for (size_t j = 0; j < true_oes[i].size(); ++j) {
       if (true_oes[i][j] != test_oes[i][j]) {
         LOG(ERROR) << "different outgoing edge";
         return false;

diff --git a/modules/graph/writer/arrow_fragment_writer.cc b/modules/graph/writer/arrow_fragment_writer.cc
@@ -65,6 +65,21 @@ void InitializeArrayArrayBuilders(
       builders[col_id] = std::make_shared<arrow::StringBuilder>();
     } else if (arrow::large_utf8()->Equals(prop_type)) {
       builders[col_id] = std::make_shared<arrow::LargeStringBuilder>();
+    } else if (arrow::date32()->Equals(prop_type)) {
+      builders[col_id] = std::make_shared<arrow::Date32Builder>();
+    } else if (arrow::date64()->Equals(prop_type)) {
+      builders[col_id] = std::make_shared<arrow::Date64Builder>();
+    } else if (prop_type->id() == arrow::Type::TIME32) {
+      builders[col_id] = std::make_shared<arrow::Time32Builder>(
+          prop_type, arrow::default_memory_pool());
+    } else if (prop_type->id() == arrow::Type::TIME64) {
+      builders[col_id] = std::make_shared<arrow::Time64Builder>(
+          prop_type, arrow::default_memory_pool());
+    } else if (prop_type->id() == arrow::Type::TIMESTAMP) {
+      builders[col_id] = std::make_shared<arrow::TimestampBuilder>(
+          prop_type, arrow::default_memory_pool());
+    } else {
+      LOG(FATAL) << "Unsupported property type: " << prop_type->ToString();
     }
     ++col_id;
   }

diff --git a/modules/graph/writer/arrow_fragment_writer_impl.h b/modules/graph/writer/arrow_fragment_writer_impl.h
@@ -456,6 +456,34 @@ ArrowFragmentWriter<FRAG_T>::appendPropertiesToArrowArrayBuilders(
           builders[col_id]);
       ARROW_OK_OR_RAISE(
           builder->Append(edge.template get_data<std::string>(pid)));
+    } else if (arrow::date32()->Equals(prop_type)) {
+      auto builder =
+          std::dynamic_pointer_cast<arrow::Date32Builder>(builders[col_id]);
+      ARROW_OK_OR_RAISE(builder->Append(
+          edge.template get_data<arrow::Date32Type::c_type>(pid)));
+    } else if (arrow::date64()->Equals(prop_type)) {
+      auto builder =
+          std::dynamic_pointer_cast<arrow::Date64Builder>(builders[col_id]);
+      ARROW_OK_OR_RAISE(builder->Append(
+          edge.template get_data<arrow::Date64Type::c_type>(pid)));
+    } else if (prop_type->id() == arrow::Type::TIME32) {
+      auto builder =
+          std::dynamic_pointer_cast<arrow::Time32Builder>(builders[col_id]);
+      ARROW_OK_OR_RAISE(builder->Append(
+          edge.template get_data<arrow::Time32Type::c_type>(pid)));
+    } else if (prop_type->id() == arrow::Type::TIME64) {
+      auto builder =
+          std::dynamic_pointer_cast<arrow::Date64Builder>(builders[col_id]);
+      ARROW_OK_OR_RAISE(builder->Append(
+          edge.template get_data<arrow::Date64Type::c_type>(pid)));
+    } else if (prop_type->id() == arrow::Type::TIMESTAMP) {
+      auto builder =
+          std::dynamic_pointer_cast<arrow::TimestampBuilder>(builders[col_id]);
+      ARROW_OK_OR_RAISE(builder->Append(
+          edge.template get_data<arrow::TimestampType::c_type>(pid)));
+    } else {
+      RETURN_GS_ERROR(ErrorCode::kInvalidValueError,
+                      "Unsupported property type: " + prop_type->ToString());
     }
     ++col_id;
   }