byte array kv meta for packed

Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Sep 10, 2024 · b6526dc · b6526dc
1 parent 3b9c278
commit b6526dc
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 11 deletions.
diff --git a/cpp/include/milvus-storage/common/fs_util.h b/cpp/include/milvus-storage/common/fs_util.h
@@ -30,4 +30,24 @@ static constexpr int64_t ARROW_PART_UPLOAD_SIZE = 10 * 1024 * 1024;  // 10 MB
 
 static constexpr int64_t MIN_BUFFER_SIZE_PER_FILE = DEFAULT_MAX_ROW_GROUP_SIZE + ARROW_PART_UPLOAD_SIZE;
 
+static const std::string ROW_GROUP_SIZE_META_KEY = "row_group_size";
+
+class PackedMetaSerde {
+  public:
+  // Serialize a vector of size_t to a byte array and convert it to a string
+  static std::string serialize(const std::vector<size_t>& sizes) {
+    std::vector<uint8_t> byteArray(sizes.size() * sizeof(size_t));
+    std::memcpy(byteArray.data(), sizes.data(), byteArray.size());
+    return std::string(byteArray.begin(), byteArray.end());
+  }
+
+  // Deserialize a string back to a vector of size_t
+  static std::vector<size_t> deserialize(const std::string& input) {
+    std::vector<uint8_t> byteArray(input.begin(), input.end());
+    std::vector<size_t> sizes(byteArray.size() / sizeof(size_t));
+    std::memcpy(sizes.data(), byteArray.data(), byteArray.size());
+    return sizes;
+  }
+};
+
 }  // namespace milvus_storage
diff --git a/cpp/src/format/parquet/file_writer.cpp b/cpp/src/format/parquet/file_writer.cpp
@@ -89,9 +89,8 @@ Status ParquetFileWriter::WriteRecordBatches(const std::vector<std::shared_ptr<a
 int64_t ParquetFileWriter::count() { return count_; }
 
 Status ParquetFileWriter::Close() {
-  std::vector<uint8_t> byteArray(row_group_sizes_.size() * sizeof(size_t));
-  std::memcpy(byteArray.data(), row_group_sizes_.data(), byteArray.size());
-  kv_metadata_->Append("row_group_sizes", std::string(byteArray.begin(), byteArray.end()));
+  std::string meta = PackedMetaSerde::serialize(row_group_sizes_);
+  kv_metadata_->Append(ROW_GROUP_SIZE_META_KEY, meta);
   RETURN_ARROW_NOT_OK(writer_->AddKeyValueMetadata(kv_metadata_));
   RETURN_ARROW_NOT_OK(writer_->Close());
   return Status::OK();

diff --git a/cpp/src/packed/reader.cpp b/cpp/src/packed/reader.cpp
@@ -19,6 +19,7 @@
 #include <parquet/properties.h>
 #include <memory>
 #include "common/arrow_util.h"
+#include "common/fs_util.h"
 #include "common/log.h"
 #include "packed/chunk_manager.h"
 
@@ -49,15 +50,9 @@ PackedRecordBatchReader::PackedRecordBatchReader(arrow::fs::FileSystem& fs,
     file_readers_.emplace_back(std::move(result.value()));
   }
 
-  auto parse_size = [&](const std::string& input) -> std::vector<size_t> {
-    std::vector<uint8_t> byteArray = std::vector<uint8_t>(input.begin(), input.end());
-    std::vector<size_t> vec(byteArray.size() / sizeof(size_t));
-    std::memcpy(vec.data(), byteArray.data(), byteArray.size());
-    return vec;
-  };
   for (int i = 0; i < file_readers_.size(); ++i) {
-    row_group_sizes_.push_back(
-        parse_size(file_readers_[i]->parquet_reader()->metadata()->key_value_metadata()->value(0)));
+    auto metadata = file_readers_[i]->parquet_reader()->metadata()->key_value_metadata()->Get(ROW_GROUP_SIZE_META_KEY);
+    row_group_sizes_.push_back(PackedMetaSerde::deserialize(metadata.ValueOrDie()));
     LOG_STORAGE_DEBUG_ << " file " << i << " metadata size: " << file_readers_[i]->parquet_reader()->metadata()->size();
   }
   // Initialize table states and chunk manager