enhance: move jsonshredding meta from parquet to meta.json (#46130)

#42533

Signed-off-by: luzhang <[email protected]>
Co-authored-by: luzhang <[email protected]>

Author: zhagnlu

internal/core/src/common/Consts.h

@@ -59,8 +59,9 @@ constexpr const char* JSON_STATS_ROOT_PATH = "json_stats";
 constexpr const char* JSON_STATS_DATA_FORMAT_VERSION = "2";
 constexpr const char* JSON_STATS_SHARED_INDEX_PATH = "shared_key_index";
 constexpr const char* JSON_STATS_SHREDDING_DATA_PATH = "shredding_data";
+constexpr const char* JSON_STATS_META_FILE_NAME = "meta.json";
 constexpr const char* JSON_KEY_STATS_SHARED_FIELD_NAME = "__shared";
-// store key layout type in parquet file metadata
+// store key layout type in parquet file metadata (deprecated, now stored in separate file)
 inline constexpr const char* JSON_STATS_META_KEY_LAYOUT_TYPE_MAP =
     "key_layout_type_map";
 // start json stats field id for mock column

internal/core/src/index/json_stats/JsonKeyStats.cpp

@@ -23,6 +23,7 @@
 #include "index/InvertedIndexUtil.h"
 #include "index/Utils.h"
 #include "milvus-storage/filesystem/fs.h"
+#include "storage/LocalChunkManagerSingleton.h"
 #include "storage/MmapManager.h"
 #include "storage/Util.h"
 #include "common/bson_view.h"
@@ -624,6 +625,64 @@ JsonKeyStats::GetSharedKeyIndexDir() {
     return shared_key_index_path.string();
 }
 
+std::string
+JsonKeyStats::GetMetaFilePath() {
+    std::filesystem::path json_stats_dir = path_;
+    std::filesystem::path meta_file_path =
+        json_stats_dir / JSON_STATS_META_FILE_NAME;
+    return meta_file_path.string();
+}
+
+void
+JsonKeyStats::WriteMetaFile() {
+    json_stats_meta_.SetLayoutTypeMap(key_types_);
+    json_stats_meta_.SetInt64(META_KEY_NUM_ROWS, num_rows_);
+    json_stats_meta_.SetInt64(META_KEY_NUM_SHREDDING_COLUMNS,
+                              column_keys_.size());
+
+    auto meta_content = json_stats_meta_.Serialize();
+    auto meta_file_path = GetMetaFilePath();
+
+    auto local_chunk_manager =
+        milvus::storage::LocalChunkManagerSingleton::GetInstance()
+            .GetChunkManager();
+    local_chunk_manager->Write(
+        meta_file_path, meta_content.data(), meta_content.size());
+
+    meta_file_size_ = meta_content.size();
+    LOG_INFO("write meta file: {} with size {} for segment {} for field {}",
+             meta_file_path,
+             meta_file_size_,
+             segment_id_,
+             field_id_);
+}
+
+void
+JsonKeyStats::LoadMetaFile(const std::string& local_meta_file_path) {
+    LOG_INFO("load meta file: {} for segment {} for field {}",
+             local_meta_file_path,
+             segment_id_,
+             field_id_);
+
+    auto local_chunk_manager =
+        storage::LocalChunkManagerSingleton::GetInstance().GetChunkManager();
+
+    auto file_size = local_chunk_manager->Size(local_meta_file_path);
+    std::string meta_content;
+    meta_content.resize(file_size);
+    local_chunk_manager->Read(
+        local_meta_file_path, meta_content.data(), file_size);
+
+    key_field_map_ = JsonStatsMeta::DeserializeToKeyFieldMap(meta_content);
+
+    LOG_INFO(
+        "loaded meta file with {} key field entries for segment {} for field "
+        "{}",
+        key_field_map_.size(),
+        segment_id_,
+        field_id_);
+}
+
 BinarySet
 JsonKeyStats::Serialize(const Config& config) {
     return BinarySet();
@@ -709,6 +768,9 @@ JsonKeyStats::BuildWithFieldData(const std::vector<FieldDataPtr>& field_datas,
     BuildKeyStats(field_datas, nullable);
     parquet_writer_->Close();
     bson_inverted_index_->BuildIndex();
+
+    // write meta file with layout type map and other metadata
+    WriteMetaFile();
 }
 
 void
@@ -725,7 +787,6 @@ JsonKeyStats::GetColumnSchemaFromParquet(int64_t column_group_id,
     LOG_DEBUG("get column schema: [{}] for segment {}",
               file_schema->ToString(true),
               segment_id_);
-    column_group_schemas_[column_group_id] = file_schema;
 
     for (const auto& field : file_schema->fields()) {
         auto field_name = field->name();
@@ -845,11 +906,18 @@ JsonKeyStats::LoadShreddingMeta(
     auto remote_prefix =
         AddBucketName(disk_file_manager_->GetRemoteJsonStatsShreddingPrefix());
 
-    // load common meta from parquet, all parquet files have the same meta
-    // just need to read one file
-    auto file = CreateColumnGroupParquetPath(
-        remote_prefix, sorted_files[0].first, sorted_files[0].second[0]);
-    GetCommonMetaFromParquet(file);
+    // load common meta from parquet only if key_field_map_ is not already populated
+    // (for backward compatibility with old data that doesn't have separate meta file)
+    if (key_field_map_.empty()) {
+        auto file = CreateColumnGroupParquetPath(
+            remote_prefix, sorted_files[0].first, sorted_files[0].second[0]);
+        GetCommonMetaFromParquet(file);
+    } else {
+        LOG_INFO(
+            "skip loading common meta from parquet, already loaded from meta "
+            "file for segment {}",
+            segment_id_);
+    }
 
     // load distinct meta from parquet, distinct meta is different for each parquet file
     // main purpose is to get column schema
@@ -1003,7 +1071,11 @@ JsonKeyStats::Load(milvus::tracer::TraceContext ctx, const Config& config) {
     AssertInfo(index_files.has_value(),
                "index file paths is empty when load json stats for segment {}",
                segment_id_);
-    // split index_files into shared_key_index and shredding_data
+
+    // split index_files into meta, shared_key_index, and shredding_data
+    // Note: Check directory paths (shared_key_index, shredding_data) BEFORE meta.json,
+    // because shared_key_index/meta.json_0 contains "meta.json" but is not the meta file.
+    std::vector<std::string> meta_files;
     std::vector<std::string> shared_key_index_files;
     std::vector<std::string> shredding_data_files;
     for (const auto& file : index_files.value()) {
@@ -1012,6 +1084,8 @@ JsonKeyStats::Load(milvus::tracer::TraceContext ctx, const Config& config) {
         } else if (file.find(JSON_STATS_SHREDDING_DATA_PATH) !=
                    std::string::npos) {
             shredding_data_files.emplace_back(file);
+        } else if (file.find(JSON_STATS_META_FILE_NAME) != std::string::npos) {
+            meta_files.emplace_back(file);
         } else {
             ThrowInfo(ErrorCode::UnexpectedError,
                       "unknown file path: {} for segment {}",
@@ -1020,6 +1094,20 @@ JsonKeyStats::Load(milvus::tracer::TraceContext ctx, const Config& config) {
         }
     }
 
+    // load meta file first (contains layout type map)
+    if (!meta_files.empty()) {
+        AssertInfo(
+            meta_files.size() == 1,
+            "expected exactly one meta file, got {} for segment {}, field {}",
+            meta_files.size(),
+            segment_id_,
+            field_id_);
+        // cache meta file to local disk
+        auto local_meta_file = disk_file_manager_->CacheJsonStatsMetaToDisk(
+            meta_files[0], load_priority_);
+        LoadMetaFile(local_meta_file);
+    }
+
     // load shredding data
     LoadShreddingData(shredding_data_files);
 
@@ -1034,15 +1122,35 @@ JsonKeyStats::Upload(const Config& config) {
     // upload inverted index
     auto bson_index_stats = bson_inverted_index_->UploadIndex();
 
+    // upload meta file
+    auto meta_file_path = GetMetaFilePath();
+    AssertInfo(disk_file_manager_->AddJsonStatsMetaLog(meta_file_path),
+               "failed to upload meta file: {} for segment {}",
+               meta_file_path,
+               segment_id_);
+
     // upload parquet file, parquet writer has already upload file to remote
     auto shredding_remote_paths_to_size = parquet_writer_->GetPathsToSize();
     auto shared_key_index_remote_paths_to_size =
         bson_index_stats->GetSerializedIndexFileInfo();
+    auto meta_remote_paths_to_size =
+        disk_file_manager_->GetRemotePathsToFileSize();
 
     // get all index files for meta
     std::vector<SerializedIndexFileInfo> index_files;
     index_files.reserve(shredding_remote_paths_to_size.size() +
-                        shared_key_index_remote_paths_to_size.size());
+                        shared_key_index_remote_paths_to_size.size() + 1);
+
+    // add meta file
+    for (const auto& [path, size] : meta_remote_paths_to_size) {
+        if (path.find(JSON_STATS_META_FILE_NAME) != std::string::npos) {
+            auto file_path = path.substr(path.find(JSON_STATS_META_FILE_NAME));
+            index_files.emplace_back(file_path, size);
+            LOG_INFO(
+                "upload meta file: {} for segment {}", file_path, segment_id_);
+        }
+    }
+
     // only store shared_key_index/... and shredding_data/... to meta
     // for saving meta space
     for (const auto& file_info : shared_key_index_remote_paths_to_size) {
@@ -1065,16 +1173,18 @@ JsonKeyStats::Upload(const Config& config) {
 
     LOG_INFO(
         "upload json key stats for segment {} with bson mem size: {} "
-        "and "
-        "shredding data mem size: {} and index files size: {}",
+        "and shredding data mem size: {} and meta file size: {} "
+        "and index files size: {}",
         segment_id_,
         bson_index_stats->GetMemSize(),
         parquet_writer_->GetTotalSize(),
+        meta_file_size_,
         index_files.size());
 
-    return IndexStats::New(
-        bson_index_stats->GetMemSize() + parquet_writer_->GetTotalSize(),
-        std::move(index_files));
+    return IndexStats::New(bson_index_stats->GetMemSize() +
+                               parquet_writer_->GetTotalSize() +
+                               meta_file_size_,
+                           std::move(index_files));
 }
 
 }  // namespace milvus::index

internal/core/src/index/json_stats/JsonKeyStats.h

@@ -457,6 +457,15 @@ class JsonKeyStats : public ScalarIndex<std::string> {
     std::string
     GetSharedKeyIndexDir();
 
+    std::string
+    GetMetaFilePath();
+
+    void
+    WriteMetaFile();
+
+    void
+    LoadMetaFile(const std::string& meta_file_path);
+
     void
     AddKeyStats(const std::vector<std::string>& path,
                 JSONType type,
@@ -648,8 +657,6 @@ class JsonKeyStats : public ScalarIndex<std::string> {
     std::unordered_map<int64_t, std::string> field_id_to_name_map_;
     // field_name vector, the sequece is the same as the order of files
     std::vector<std::string> field_names_;
-    // column_group_id -> schema, the sequence of schemas is the same as the order of files
-    std::map<int64_t, std::shared_ptr<arrow::Schema>> column_group_schemas_;
     // field_name -> column
     mutable std::unordered_map<std::string,
                                std::shared_ptr<milvus::ChunkedColumnInterface>>
@@ -661,6 +668,10 @@ class JsonKeyStats : public ScalarIndex<std::string> {
     SkipIndex skip_index_;
     cachinglayer::ResourceUsage cell_size_ = {0, 0};
 
+    // Meta file for storing layout type map and other metadata
+    JsonStatsMeta json_stats_meta_;
+    int64_t meta_file_size_{0};
+
     // Friend accessor for unit tests to call private methods safely.
     friend class ::TraverseJsonForBuildStatsAccessor;
     friend class ::CollectSingleJsonStatsInfoAccessor;

internal/core/src/index/json_stats/utils.cpp

@@ -18,6 +18,7 @@
 #include <boost/filesystem.hpp>
 #include <boost/uuid/random_generator.hpp>
 #include <boost/uuid/uuid_io.hpp>
+#include <nlohmann/json.hpp>
 #include "milvus-storage/common/constants.h"
 
 namespace milvus::index {
@@ -207,15 +208,101 @@ CreateArrowSchema(std::map<JsonKey, JsonKeyLayoutType> column_map) {
 
 std::vector<std::pair<std::string, std::string>>
 CreateParquetKVMetadata(std::map<JsonKey, JsonKeyLayoutType> column_map) {
-    nlohmann::json key_type;
-    for (const auto& [key, type] : column_map) {
-        key_type[key.ToColumnName()] = ToString(type);
+    // layout type map is now stored in a separate meta file to reduce parquet file size.
+    // return empty metadata vector.
+    return {};
+}
+
+std::string
+JsonStatsMeta::Serialize() const {
+    nlohmann::json root;
+
+    // Serialize string values
+    for (const auto& [key, value] : string_values_) {
+        root[key] = value;
     }
-    // for shared field, not need to store in metadata
-    std::vector<std::pair<std::string, std::string>> res;
-    res.push_back(
-        std::make_pair(JSON_STATS_META_KEY_LAYOUT_TYPE_MAP, key_type.dump()));
-    return res;
+
+    // Serialize int64 values
+    for (const auto& [key, value] : int64_values_) {
+        root[key] = value;
+    }
+
+    // Serialize layout type map
+    if (!layout_type_map_.empty()) {
+        nlohmann::json layout_map;
+        for (const auto& [json_key, layout_type] : layout_type_map_) {
+            layout_map[json_key.ToColumnName()] = ToString(layout_type);
+        }
+        root[META_KEY_LAYOUT_TYPE_MAP] = layout_map;
+    }
+
+    return root.dump();
+}
+
+JsonStatsMeta
+JsonStatsMeta::Deserialize(const std::string& json_str) {
+    JsonStatsMeta meta;
+
+    try {
+        nlohmann::json root = nlohmann::json::parse(json_str);
+
+        for (auto it = root.begin(); it != root.end(); ++it) {
+            const std::string& key = it.key();
+
+            if (key == META_KEY_LAYOUT_TYPE_MAP) {
+                // Parse layout type map
+                std::map<JsonKey, JsonKeyLayoutType> layout_map;
+                for (auto& [column_name, layout_type_str] :
+                     it.value().items()) {
+                    auto json_type = GetJsonTypeFromKeyName(column_name);
+                    auto json_pointer = GetKeyFromColumnName(column_name);
+                    JsonKey json_key(json_pointer, json_type);
+                    auto layout_type =
+                        JsonKeyLayoutTypeFromString(layout_type_str);
+                    layout_map[json_key] = layout_type;
+                }
+                meta.SetLayoutTypeMap(layout_map);
+            } else if (it.value().is_string()) {
+                meta.SetString(key, it.value().get<std::string>());
+            } else if (it.value().is_number_integer()) {
+                meta.SetInt64(key, it.value().get<int64_t>());
+            }
+            // Other types can be added as needed
+        }
+    } catch (const std::exception& e) {
+        ThrowInfo(ErrorCode::UnexpectedError,
+                  "Failed to deserialize JsonStatsMeta: {}",
+                  e.what());
+    }
+
+    return meta;
+}
+
+std::unordered_map<std::string, std::set<std::string>>
+JsonStatsMeta::DeserializeToKeyFieldMap(const std::string& json_str) {
+    std::unordered_map<std::string, std::set<std::string>> key_field_map;
+
+    try {
+        nlohmann::json root = nlohmann::json::parse(json_str);
+
+        auto it = root.find(META_KEY_LAYOUT_TYPE_MAP);
+        if (it != root.end()) {
+            for (auto& [column_name, layout_type_str] : it.value().items()) {
+                auto layout_type = JsonKeyLayoutTypeFromString(layout_type_str);
+                if (layout_type == JsonKeyLayoutType::SHARED) {
+                    continue;
+                }
+                auto json_pointer = GetKeyFromColumnName(column_name);
+                key_field_map[json_pointer].insert(column_name);
+            }
+        }
+    } catch (const std::exception& e) {
+        ThrowInfo(ErrorCode::UnexpectedError,
+                  "Failed to deserialize JsonStatsMeta to key_field_map: {}",
+                  e.what());
+    }
+
+    return key_field_map;
 }
 
 }  // namespace milvus::index