1 | #pragma once |
2 | #include <functional> |
3 | #include <IO/WriteBuffer.h> |
4 | |
5 | |
6 | namespace DB |
7 | { |
8 | |
9 | namespace ErrorCodes |
10 | { |
11 | extern const int CURRENT_WRITE_BUFFER_IS_EXHAUSTED; |
12 | } |
13 | |
14 | /* The buffer is similar to ConcatReadBuffer, but writes data |
15 | * |
16 | * It has WriteBuffers sequence [prepared_sources, lazy_sources] |
17 | * (lazy_sources contains not pointers themself, but their delayed constructors) |
18 | * |
19 | * Firtly, CascadeWriteBuffer redirects data to first buffer of the sequence |
20 | * If current WriteBuffer cannot receive data anymore, it throws special exception CURRENT_WRITE_BUFFER_IS_EXHAUSTED in nextImpl() body, |
21 | * CascadeWriteBuffer prepare next buffer and continuously redirects data to it. |
22 | * If there are no buffers anymore CascadeWriteBuffer throws an exception. |
23 | * |
24 | * NOTE: If you use one of underlying WriteBuffers buffers outside, you need sync its position() with CascadeWriteBuffer's position(). |
25 | * The sync is performed into nextImpl(), getResultBuffers() and destructor. |
26 | */ |
27 | class CascadeWriteBuffer : public WriteBuffer |
28 | { |
29 | public: |
30 | |
31 | using WriteBufferPtrs = std::vector<WriteBufferPtr>; |
32 | using WriteBufferConstructor = std::function<WriteBufferPtr (const WriteBufferPtr & prev_buf)>; |
33 | using WriteBufferConstructors = std::vector<WriteBufferConstructor>; |
34 | |
35 | CascadeWriteBuffer(WriteBufferPtrs && prepared_sources_, WriteBufferConstructors && lazy_sources_ = {}); |
36 | |
37 | void nextImpl() override; |
38 | |
39 | /// Should be called once |
40 | void getResultBuffers(WriteBufferPtrs & res); |
41 | |
42 | const WriteBuffer * getCurrentBuffer() const |
43 | { |
44 | return curr_buffer; |
45 | } |
46 | |
47 | ~CascadeWriteBuffer() override; |
48 | |
49 | private: |
50 | |
51 | WriteBuffer * setNextBuffer(); |
52 | |
53 | WriteBufferPtrs prepared_sources; |
54 | WriteBufferConstructors lazy_sources; |
55 | size_t first_lazy_source_num; |
56 | size_t num_sources; |
57 | |
58 | WriteBuffer * curr_buffer; |
59 | size_t curr_buffer_num; |
60 | }; |
61 | |
62 | } |
63 | |