Less copy between in-process memory and vineyard during building a fr…

…agment (#1154)

Signed-off-by: Tao He <[email protected]>

modules/basic/ds/arrow.cc

@@ -314,6 +314,59 @@ template class NumericArrayBuilder<uint64_t>;
 template class NumericArrayBuilder<float>;
 template class NumericArrayBuilder<double>;
 
+template <typename T>
+FixedNumericArrayBuilder<T>::FixedNumericArrayBuilder(Client& client,
+                                                      const size_t size)
+    : NumericArrayBaseBuilder<T>(client), size_(size) {
+  if (size_ > 0) {
+    VINEYARD_CHECK_OK(client.CreateBlob(size_ * sizeof(T), writer_));
+    data_ = reinterpret_cast<T*>(writer_->data());
+  }
+}
+
+template <typename T>
+size_t FixedNumericArrayBuilder<T>::size() const {
+  return size_;
+}
+
+template <typename T>
+T* FixedNumericArrayBuilder<T>::MutablePointer(int64_t i) const {
+  if (data_) {
+    return data_ + i;
+  }
+  return nullptr;
+}
+
+template <typename T>
+T* FixedNumericArrayBuilder<T>::data() const {
+  return data_;
+}
+
+template <typename T>
+Status FixedNumericArrayBuilder<T>::Build(Client& client) {
+  this->set_length_(size_);
+  this->set_null_count_(0);
+  this->set_offset_(0);
+  if (size_ > 0) {
+    this->set_buffer_(std::move(writer_));
+  } else {
+    this->set_buffer_(Blob::MakeEmpty(client));
+  }
+  this->set_null_bitmap_(Blob::MakeEmpty(client));
+  return Status::OK();
+}
+
+template class FixedNumericArrayBuilder<int8_t>;
+template class FixedNumericArrayBuilder<int16_t>;
+template class FixedNumericArrayBuilder<int32_t>;
+template class FixedNumericArrayBuilder<int64_t>;
+template class FixedNumericArrayBuilder<uint8_t>;
+template class FixedNumericArrayBuilder<uint16_t>;
+template class FixedNumericArrayBuilder<uint32_t>;
+template class FixedNumericArrayBuilder<uint64_t>;
+template class FixedNumericArrayBuilder<float>;
+template class FixedNumericArrayBuilder<double>;
+
 BooleanArrayBuilder::BooleanArrayBuilder(Client& client)
     : BooleanArrayBaseBuilder(client) {
   std::shared_ptr<ArrayType> array;

modules/basic/ds/arrow.h

@@ -77,6 +77,34 @@ class NumericArrayBuilder : public NumericArrayBaseBuilder<T> {
   std::vector<std::shared_ptr<arrow::Array>> arrays_;
 };
 
+/**
+ * @brief FixedNumericArrayBuilder is designed for building Arrow numeric
+ * arrays with known size. It is useful for allocating buffers directly
+ * on the vineyard's shared memory.
+ *
+ * @tparam T
+ */
+template <typename T>
+class FixedNumericArrayBuilder : public NumericArrayBaseBuilder<T> {
+ public:
+  using ArrayType = ArrowArrayType<T>;
+
+  FixedNumericArrayBuilder(Client& client, const size_t size);
+
+  size_t size() const;
+
+  T* MutablePointer(int64_t i) const;
+
+  T* data() const;
+
+  Status Build(Client& client) override;
+
+ private:
+  size_t size_ = 0;
+  std::unique_ptr<BlobWriter> writer_ = nullptr;
+  T* data_ = nullptr;
+};
+
 using Int8Builder = NumericArrayBuilder<int8_t>;
 using Int16Builder = NumericArrayBuilder<int16_t>;
 using Int32Builder = NumericArrayBuilder<int32_t>;
@@ -88,6 +116,17 @@ using UInt64Builder = NumericArrayBuilder<uint64_t>;
 using FloatBuilder = NumericArrayBuilder<float>;
 using DoubleBuilder = NumericArrayBuilder<double>;
 
+using FixedInt8Builder = FixedNumericArrayBuilder<int8_t>;
+using FixedInt16Builder = FixedNumericArrayBuilder<int16_t>;
+using FixedInt32Builder = FixedNumericArrayBuilder<int32_t>;
+using FixedInt64Builder = FixedNumericArrayBuilder<int64_t>;
+using FixedUInt8Builder = FixedNumericArrayBuilder<uint8_t>;
+using FixedUInt16Builder = FixedNumericArrayBuilder<uint16_t>;
+using FixedUInt32Builder = FixedNumericArrayBuilder<uint32_t>;
+using FixedUInt64Builder = FixedNumericArrayBuilder<uint64_t>;
+using FixedFloatBuilder = FixedNumericArrayBuilder<float>;
+using FixedDoubleBuilder = FixedNumericArrayBuilder<double>;
+
 /**
  * @brief BooleanArrayBuilder is designed for constructing  Arrow arrays of
  * boolean data type
@@ -194,13 +233,15 @@ class PodArrayBuilder : public FixedSizeBinaryArrayBaseBuilder {
     }
   }
 
-  T* MutablePointer(int64_t i) {
+  T* MutablePointer(int64_t i) const {
     if (data_) {
       return data_ + i;
     }
     return nullptr;
   }
 
+  T* data() const { return data_; }
+
   size_t size() const { return size_; }
 
   Status Build(Client& client) override {

modules/basic/ds/arrow_utils.h

@@ -50,6 +50,9 @@ class NumericArray;
 template <typename T>
 class NumericArrayBuilder;
 
+template <typename T>
+class FixedNumericArrayBuilder;
+
 #define CONVERT_TO_ARROW_TYPE(type, data_type, array_type, builder_type,       \
                               type_value)                                      \
   template <>                                                                  \
@@ -59,6 +62,7 @@ class NumericArrayBuilder;
     using VineyardArrayType = NumericArray<type>;                              \
     using BuilderType = builder_type;                                          \
     using VineyardBuilderType = NumericArrayBuilder<type>;                     \
+    using FixedVineyardBuilderType = FixedNumericArrayBuilder<type>;           \
     static std::shared_ptr<arrow::DataType> TypeValue() { return type_value; } \
   };
 
@@ -79,6 +83,10 @@ template <typename T>
 using ArrowVineyardBuilderType =
     typename ConvertToArrowType<T>::VineyardBuilderType;
 
+template <typename T>
+using FixedArrowVineyardBuilderType =
+    typename ConvertToArrowType<T>::FixedVineyardBuilderType;
+
 CONVERT_TO_ARROW_TYPE(void, arrow::NullType, arrow::NullArray,
                       arrow::NullBuilder, arrow::null())
 CONVERT_TO_ARROW_TYPE(bool, arrow::BooleanType, arrow::BooleanArray,

modules/graph/fragment/arrow_fragment.h

@@ -112,7 +112,7 @@ class BasicArrowFragmentBuilder
 
   std::vector<std::vector<std::shared_ptr<PodArrayBuilder<nbr_unit_t>>>>
       ie_lists_, oe_lists_;
-  std::vector<std::vector<std::shared_ptr<arrow::Int64Array>>>
+  std::vector<std::vector<std::shared_ptr<FixedInt64Builder>>>
       ie_offsets_lists_, oe_offsets_lists_;
 
   std::shared_ptr<vertex_map_t> vm_ptr_;

modules/graph/fragment/arrow_fragment.vineyard-mod

@@ -678,7 +678,7 @@ class [[vineyard]] ArrowFragment
       vineyard::Client & client,
       std::vector<std::vector<std::shared_ptr<PodArrayBuilder<nbr_unit_t>>>> &
           oe_lists,
-      std::vector<std::vector<std::shared_ptr<arrow::Int64Array>>> &
+      std::vector<std::vector<std::shared_ptr<FixedInt64Builder>>> &
           oe_offsets_lists,
       int concurrency, bool& is_multigraph);