forked from alibaba/AliSQL
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathwindow_collection.cpp
More file actions
146 lines (122 loc) · 4.16 KB
/
window_collection.cpp
File metadata and controls
146 lines (122 loc) · 4.16 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
#include "duckdb/function/window/window_collection.hpp"
namespace duckdb {
//===--------------------------------------------------------------------===//
// WindowCollection
//===--------------------------------------------------------------------===//
WindowCollection::WindowCollection(BufferManager &buffer_manager, idx_t count, const vector<LogicalType> &types)
: all_valids(types.size()), types(types), count(count), buffer_manager(buffer_manager) {
if (!types.empty()) {
inputs = make_uniq<ColumnDataCollection>(buffer_manager, types);
}
validities.resize(types.size());
// Atomic vectors can't be constructed with a given value
for (auto &all_valid : all_valids) {
all_valid = true;
}
}
void WindowCollection::GetCollection(idx_t row_idx, ColumnDataCollectionSpec &spec) {
if (spec.second && row_idx == spec.first + spec.second->Count()) {
return;
}
lock_guard<mutex> collection_guard(lock);
auto collection = make_uniq<ColumnDataCollection>(buffer_manager, types);
spec = {row_idx, collection.get()};
Range probe {row_idx, collections.size()};
auto i = std::upper_bound(ranges.begin(), ranges.end(), probe);
ranges.insert(i, probe);
collections.emplace_back(std::move(collection));
}
void WindowCollection::Combine(const ColumnSet &validity_cols) {
lock_guard<mutex> collection_guard(lock);
// If there are no columns (COUNT(*)) then this is a NOP
if (types.empty()) {
return;
}
// Have we already combined?
if (inputs->Count()) {
D_ASSERT(collections.empty());
D_ASSERT(ranges.empty());
return;
}
// If there are columns, we should have data
D_ASSERT(!collections.empty());
D_ASSERT(!ranges.empty());
for (auto &range : ranges) {
inputs->Combine(*collections[range.second]);
}
collections.clear();
ranges.clear();
if (validity_cols.empty()) {
return;
}
D_ASSERT(inputs.get());
// Find all columns with NULLs
vector<column_t> invalid_cols;
for (auto &col_idx : validity_cols) {
if (!all_valids[col_idx]) {
invalid_cols.emplace_back(col_idx);
validities[col_idx].Initialize(inputs->Count());
}
}
if (invalid_cols.empty()) {
return;
}
WindowCursor cursor(*this, invalid_cols);
idx_t target_offset = 0;
while (cursor.Scan()) {
const auto count = cursor.chunk.size();
for (idx_t i = 0; i < invalid_cols.size(); ++i) {
auto &other = FlatVector::Validity(cursor.chunk.data[i]);
const auto col_idx = invalid_cols[i];
validities[col_idx].SliceInPlace(other, target_offset, 0, count);
}
target_offset += count;
}
}
WindowBuilder::WindowBuilder(WindowCollection &collection) : collection(collection) {
}
void WindowBuilder::Sink(DataChunk &chunk, idx_t input_idx) {
// Check whether we need a a new collection
if (!sink.second || input_idx < sink.first || sink.first + sink.second->Count() < input_idx) {
collection.GetCollection(input_idx, sink);
D_ASSERT(sink.second);
sink.second->InitializeAppend(appender);
}
sink.second->Append(appender, chunk);
// Record NULLs
for (column_t col_idx = 0; col_idx < chunk.ColumnCount(); ++col_idx) {
if (!collection.all_valids[col_idx]) {
continue;
}
// Column was valid, make sure it still is.
UnifiedVectorFormat data;
chunk.data[col_idx].ToUnifiedFormat(chunk.size(), data);
if (!data.validity.AllValid()) {
collection.all_valids[col_idx] = false;
}
}
}
WindowCursor::WindowCursor(const WindowCollection &paged, vector<column_t> column_ids) : paged(paged) {
D_ASSERT(paged.collections.empty());
D_ASSERT(paged.ranges.empty());
if (column_ids.empty()) {
// For things like COUNT(*) set the state up to contain the whole range
state.segment_index = 0;
state.chunk_index = 0;
state.current_row_index = 0;
state.next_row_index = paged.size();
state.properties = ColumnDataScanProperties::ALLOW_ZERO_COPY;
chunk.SetCapacity(state.next_row_index);
chunk.SetCardinality(state.next_row_index);
return;
} else if (chunk.data.empty()) {
auto &inputs = paged.inputs;
D_ASSERT(inputs.get());
inputs->InitializeScan(state, std::move(column_ids));
inputs->InitializeScanChunk(state, chunk);
}
}
WindowCursor::WindowCursor(const WindowCollection &paged, column_t col_idx)
: WindowCursor(paged, vector<column_t>(1, col_idx)) {
}
} // namespace duckdb