1#include "duckdb/execution/operator/helper/physical_limit.hpp"
2
3using namespace duckdb;
4using namespace std;
5
6class PhysicalLimitOperatorState : public PhysicalOperatorState {
7public:
8 PhysicalLimitOperatorState(PhysicalOperator *child, idx_t current_offset = 0)
9 : PhysicalOperatorState(child), current_offset(current_offset) {
10 }
11
12 idx_t current_offset;
13};
14
15void PhysicalLimit::GetChunkInternal(ClientContext &context, DataChunk &chunk, PhysicalOperatorState *state_) {
16 auto state = reinterpret_cast<PhysicalLimitOperatorState *>(state_);
17
18 idx_t max_element = limit + offset;
19 if (state->current_offset >= max_element) {
20 return;
21 }
22
23 // get the next chunk from the child
24 do {
25 children[0]->GetChunk(context, state->child_chunk, state->child_state.get());
26 if (state->child_chunk.size() == 0) {
27 return;
28 }
29
30 if (state->current_offset < offset) {
31 // we are not yet at the offset point
32 if (state->current_offset + state->child_chunk.size() > offset) {
33 // however we will reach it in this chunk
34 // we have to copy part of the chunk with an offset
35 idx_t start_position = offset - state->current_offset;
36 idx_t chunk_count = min(limit, state->child_chunk.size() - start_position);
37 SelectionVector sel(STANDARD_VECTOR_SIZE);
38 for (idx_t i = 0; i < chunk_count; i++) {
39 sel.set_index(i, start_position + i);
40 }
41 // set up a slice of the input chunks
42 chunk.Slice(state->child_chunk, sel, chunk_count);
43 }
44 } else {
45 // have to copy either the entire chunk or part of it
46 idx_t chunk_count;
47 if (state->current_offset + state->child_chunk.size() >= max_element) {
48 // have to limit the count of the chunk
49 chunk_count = max_element - state->current_offset;
50 } else {
51 // we copy the entire chunk
52 chunk_count = state->child_chunk.size();
53 }
54 // instead of copying we just change the pointer in the current chunk
55 chunk.Reference(state->child_chunk);
56 chunk.SetCardinality(chunk_count);
57 }
58
59 state->current_offset += state->child_chunk.size();
60 } while (chunk.size() == 0);
61}
62
63unique_ptr<PhysicalOperatorState> PhysicalLimit::GetOperatorState() {
64 return make_unique<PhysicalLimitOperatorState>(children[0].get(), 0);
65}
66