1 | //===----------------------------------------------------------------------===// |
---|---|
2 | // |
3 | // DuckDB |
4 | // |
5 | // duckdb_benchmark.hpp |
6 | // |
7 | // Author: Mark Raasveldt |
8 | // |
9 | //===----------------------------------------------------------------------===// |
10 | |
11 | #pragma once |
12 | |
13 | #include "benchmark.hpp" |
14 | #include "duckdb.hpp" |
15 | #include "duckdb/main/client_context.hpp" |
16 | #include "test_helpers.hpp" |
17 | |
18 | namespace duckdb { |
19 | |
20 | //! Base class for any state that has to be kept by a Benchmark |
21 | struct DuckDBBenchmarkState : public BenchmarkState { |
22 | DuckDB db; |
23 | Connection conn; |
24 | unique_ptr<QueryResult> result; |
25 | |
26 | DuckDBBenchmarkState(string path) : db(path.empty() ? nullptr : path.c_str()), conn(db) { |
27 | conn.EnableProfiling(); |
28 | } |
29 | virtual ~DuckDBBenchmarkState() { |
30 | } |
31 | }; |
32 | |
33 | //! The base Benchmark class is a base class that is used to create and register |
34 | //! new benchmarks |
35 | class DuckDBBenchmark : public Benchmark { |
36 | public: |
37 | DuckDBBenchmark(bool register_benchmark, string name, string group) : Benchmark(register_benchmark, name, group) { |
38 | } |
39 | virtual ~DuckDBBenchmark() { |
40 | } |
41 | |
42 | //! Load data into DuckDB |
43 | virtual void Load(DuckDBBenchmarkState *state) = 0; |
44 | //! A single query to run against the database |
45 | virtual string GetQuery() { |
46 | return string(); |
47 | } |
48 | //! Run a bunch of queries, only called if GetQuery() returns an empty string |
49 | virtual void RunBenchmark(DuckDBBenchmarkState *state) { |
50 | } |
51 | //! This function gets called after the GetQuery() method |
52 | virtual void Cleanup(DuckDBBenchmarkState *state){}; |
53 | //! Verify a result |
54 | virtual string VerifyResult(QueryResult *result) = 0; |
55 | //! Whether or not the benchmark is performed on an in-memory database |
56 | virtual bool InMemory() { |
57 | return true; |
58 | } |
59 | |
60 | string GetDatabasePath() { |
61 | if (!InMemory()) { |
62 | string path = "duckdb_benchmark_db.db"; |
63 | DeleteDatabase(path); |
64 | return path; |
65 | } else { |
66 | return string(); |
67 | } |
68 | } |
69 | |
70 | virtual unique_ptr<DuckDBBenchmarkState> CreateBenchmarkState() { |
71 | return make_unique<DuckDBBenchmarkState>(GetDatabasePath()); |
72 | } |
73 | |
74 | unique_ptr<BenchmarkState> Initialize() override { |
75 | auto state = CreateBenchmarkState(); |
76 | Load(state.get()); |
77 | return move(state); |
78 | } |
79 | |
80 | void Run(BenchmarkState *state_) override { |
81 | auto state = (DuckDBBenchmarkState *)state_; |
82 | string query = GetQuery(); |
83 | if (query.empty()) { |
84 | RunBenchmark(state); |
85 | } else { |
86 | state->result = state->conn.Query(query); |
87 | } |
88 | } |
89 | |
90 | void Cleanup(BenchmarkState *state_) override { |
91 | auto state = (DuckDBBenchmarkState *)state_; |
92 | Cleanup(state); |
93 | } |
94 | |
95 | string Verify(BenchmarkState *state_) override { |
96 | auto state = (DuckDBBenchmarkState *)state_; |
97 | return VerifyResult(state->result.get()); |
98 | } |
99 | |
100 | string GetLogOutput(BenchmarkState *state_) override { |
101 | auto state = (DuckDBBenchmarkState *)state_; |
102 | return state->conn.context->profiler.ToJSON(); |
103 | } |
104 | |
105 | //! Interrupt the benchmark because of a timeout |
106 | void Interrupt(BenchmarkState *state_) override { |
107 | auto state = (DuckDBBenchmarkState *)state_; |
108 | state->conn.Interrupt(); |
109 | } |
110 | }; |
111 | |
112 | } // namespace duckdb |
113 |