Write pond to file

searchlib/src/tests/queryeval/iterator_benchmark/CMakeLists.txt

@@ -5,6 +5,7 @@ vespa_add_executable(searchlib_iterator_benchmark_test_app TEST
     benchmark_blueprint_factory.cpp
     blueprint_factory_builder.cpp
     common.cpp
+    data_pond_utils.cpp
     disk_index_builder.cpp
     intermediate_blueprint_factory.cpp
     iterator_benchmark_test.cpp

searchlib/src/tests/queryeval/iterator_benchmark/data_pond_utils.cpp

@@ -0,0 +1,104 @@
+// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "data_pond_utils.h"
+
+#include <vespa/vespalib/data/memory.h>
+#include <vespa/vespalib/data/simple_buffer.h>
+#include <vespa/vespalib/data/slime/binary_format.h>
+#include <vespa/vespalib/data/slime/inspector.h>
+#include <vespa/vespalib/data/slime/object_traverser.h>
+#include <vespa/vespalib/data/slime/slime.h>
+#include <vespa/vespalib/data/slime/type.h>
+#include <vespa/vespalib/io/fileutil.h>
+
+using vespalib::File;
+using vespalib::Memory;
+using vespalib::SimpleBuffer;
+using vespalib::Slime;
+using vespalib::slime::BinaryFormat;
+using vespalib::slime::BOOL;
+using vespalib::slime::Cursor;
+using vespalib::slime::DOUBLE;
+using vespalib::slime::Inspector;
+using vespalib::slime::LONG;
+using vespalib::slime::ObjectTraverser;
+using vespalib::slime::STRING;
+
+namespace search::queryeval::test {
+
+void write_data_pond_to_file(const std::string& file_path, const DataPond& data_pond) {
+    Slime   slime;
+    Cursor& root = slime.setObject();
+    Cursor& records = root.setArray("records", data_pond.records().size());
+    for (const auto& rec : data_pond.records()) {
+        Cursor& obj = records.addObject();
+        for (const auto& [name, field] : rec.data()) {
+            if (field.has_type<int64_t>()) {
+                obj.setLong(name, field.get<int64_t>());
+            } else if (field.has_type<double>()) {
+                obj.setDouble(name, field.get<double>());
+            } else if (field.has_type<bool>()) {
+                obj.setBool(name, field.get<bool>());
+            } else if (field.has_type<std::string>()) {
+                obj.setString(name, field.get<std::string>());
+            }
+        }
+    }
+
+    SimpleBuffer buffer;
+    BinaryFormat::encode(slime, buffer);
+    Memory bytes = buffer.get();
+
+    File out(file_path);
+    out.open(File::CREATE | File::TRUNC);
+    out.write(bytes.data, bytes.size, 0);
+}
+
+/**
+ * Copies fields from slime to data pond record.
+ */
+class FieldCopier : public ObjectTraverser {
+    Record& _target;
+
+public:
+    explicit FieldCopier(Record& target) noexcept : _target(target) {}
+
+    void field(const Memory& name, const Inspector& value) override {
+        std::string field_name(name.data, name.size);
+        switch (value.type().getId()) {
+        case LONG::ID:
+            _target.set(field_name, value.asLong());
+            break;
+        case DOUBLE::ID:
+            _target.set(field_name, value.asDouble());
+            break;
+        case BOOL::ID:
+            _target.set(field_name, value.asBool());
+            break;
+        case STRING::ID: {
+            Memory s = value.asString();
+            _target.set(field_name, std::string(s.data, s.size));
+            break;
+        }
+        default:
+            break;
+        }
+    }
+};
+
+void read_file_into_data_pond(const std::string& file_path, DataPond& data_pond) {
+    std::string bytes = File::readAll(file_path);
+
+    Slime slime;
+    BinaryFormat::decode(Memory(bytes.data(), bytes.size()), slime);
+
+    const Inspector& records = slime.get()["records"];
+    for (size_t i = 0; i < records.entries(); ++i) {
+        const Inspector& rec_in = records[i];
+        Record&          rec_out = data_pond.new_record();
+        FieldCopier      copier(rec_out);
+        rec_in.traverse(static_cast<ObjectTraverser&>(copier));
+    }
+}
+
+} // namespace search::queryeval::test

searchlib/src/tests/queryeval/iterator_benchmark/data_pond_utils.h

@@ -0,0 +1,22 @@
+// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "data_pond.h"
+
+#include <string>
+
+namespace search::queryeval::test {
+
+/**
+ * Writes the records of a data pond to file.
+ */
+void write_data_pond_to_file(const std::string& file_path, const DataPond& data_pond);
+
+/**
+ * Reads records from file and adds them to the data pond.
+ */
+void read_file_into_data_pond(const std::string& file_path, DataPond& data_pond);
+
+}
+

searchlib/src/tests/queryeval/iterator_benchmark/iterator_benchmark_test.cpp

@@ -4,19 +4,23 @@
 #include "blueprint_factory_builder.h"
 #include "common.h"
 #include "data_pond.h"
+#include "data_pond_utils.h"
 #include "intermediate_blueprint_factory.h"
 
 #include <vespa/searchlib/fef/matchdata.h>
 #include <vespa/searchlib/queryeval/blueprint.h>
 #include <vespa/vespalib/gtest/gtest.h>
 #include <vespa/vespalib/util/benchmark_timer.h>
+#include <vespa/vespalib/util/exception.h>
 #include <vespa/vespalib/util/stringfmt.h>
 
 #include <cmath>
+#include <exception>
 #include <format>
 #include <functional>
 #include <iomanip>
 #include <numeric>
+#include <optional>
 #include <print>
 #include <vector>
 
@@ -1197,7 +1201,9 @@ static std::string smoke_test_filter = "--gtest_filter="
                                        "IteratorBenchmark.analyze_ENN";
 
 int main(int argc, char** argv) {
-    bool opt_dump_pond = false;
+    bool                       opt_dump_pond = false;
+    std::optional<std::string> opt_save_pond = std::nullopt;
+    std::optional<std::string> opt_load_pond = std::nullopt;
     for (int i = 0; i < argc; i++) {
         const std::string& smoke_test{"--smoke-test"};
         if (smoke_test == argv[i]) {
@@ -1208,15 +1214,52 @@ int main(int argc, char** argv) {
         if (dump_pond_flag == argv[i]) {
             opt_dump_pond = true;
         }
+        const std::string& save_pond_flag{"--save-pond"};
+        if (save_pond_flag == argv[i]) {
+            if (i + 1 >= argc) {
+                std::println(stderr, "Expected --save-pond <FILE>, but got no argument");
+                return 1;
+            }
+            opt_save_pond = std::string(argv[++i]);
+            continue;
+        }
+        const std::string& load_pond_flag{"--load-pond"};
+        if (load_pond_flag == argv[i]) {
+            if (i + 1 >= argc) {
+                std::println(stderr, "Expected --load-pond <FILE>, but got no argument");
+                return 1;
+            }
+            opt_load_pond = std::string(argv[++i]);
+            continue;
+        }
+    }
+    if (opt_load_pond) {
+        try {
+            read_file_into_data_pond(*opt_load_pond, global_pond);
+        } catch (const vespalib::Exception& e) {
+            std::println(stderr, "error: failed to load pond from '{}': {}", *opt_load_pond, e.getMessage());
+            return 1;
+        }
+    } else {
+        ::testing::InitGoogleTest(&argc, argv);
+        int res = RUN_ALL_TESTS();
+        if (res != 0) {
+            return res;
+        }
+        if (opt_save_pond) {
+            try {
+                write_data_pond_to_file(*opt_save_pond, global_pond);
+            } catch (const vespalib::Exception& e) {
+                std::println(stderr, "error: failed to save pond to '{}': {}", *opt_save_pond, e.getMessage());
+                return 1;
+            }
+        }
     }
-    ::testing::InitGoogleTest(&argc, argv);
-    int res = RUN_ALL_TESTS();
     if (!global_pond.records().empty()) {
         postprocess_pond(global_pond);
         print_pond_summary(global_pond);
     }
     if (opt_dump_pond) {
         dump_pond(global_pond);
     }
-    return res;
 }