Coverage Report

Created: 2026-05-19 11:20

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
be/src/storage/segment/column_writer.cpp
Line
Count
Source
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
#include "storage/segment/column_writer.h"
19
20
#include <gen_cpp/segment_v2.pb.h>
21
22
#include <algorithm>
23
#include <cstring>
24
#include <filesystem>
25
#include <memory>
26
27
#include "common/config.h"
28
#include "common/logging.h"
29
#include "core/data_type/data_type_agg_state.h"
30
#include "core/data_type/data_type_factory.hpp"
31
#include "core/types.h"
32
#include "io/fs/file_writer.h"
33
#include "runtime/collection_value.h"
34
#include "storage/field.h"
35
#include "storage/index/bloom_filter/bloom_filter_index_writer.h"
36
#include "storage/index/inverted/inverted_index_writer.h"
37
#include "storage/index/ordinal_page_index.h"
38
#include "storage/index/zone_map/zone_map_index.h"
39
#include "storage/olap_common.h"
40
#include "storage/segment/encoding_info.h"
41
#include "storage/segment/options.h"
42
#include "storage/segment/page_builder.h"
43
#include "storage/segment/page_io.h"
44
#include "storage/segment/page_pointer.h"
45
#include "storage/segment/variant/variant_column_writer_impl.h"
46
#include "storage/tablet/tablet_schema.h"
47
#include "storage/types.h"
48
#include "util/block_compression.h"
49
#include "util/debug_points.h"
50
#include "util/faststring.h"
51
#include "util/rle_encoding.h"
52
#include "util/simd/bits.h"
53
54
namespace doris::segment_v2 {
55
56
class NullBitmapBuilder {
57
public:
58
73.0k
    NullBitmapBuilder() : _has_null(false), _bitmap_buf(512), _rle_encoder(&_bitmap_buf, 1) {}
59
60
    explicit NullBitmapBuilder(size_t reserve_bits)
61
            : _has_null(false),
62
              _bitmap_buf(BitmapSize(reserve_bits)),
63
0
              _rle_encoder(&_bitmap_buf, 1) {}
64
65
72.9k
    void reserve_for_write(size_t num_rows, size_t non_null_count) {
66
72.9k
        if (num_rows == 0) {
67
0
            return;
68
0
        }
69
72.9k
        if (non_null_count == 0 || (non_null_count == num_rows && !_has_null)) {
70
67.8k
            if (_bitmap_buf.capacity() < kSmallReserveBytes) {
71
80
                _bitmap_buf.reserve(kSmallReserveBytes);
72
80
            }
73
67.8k
            return;
74
67.8k
        }
75
5.11k
        size_t raw_bytes = BitmapSize(num_rows);
76
5.11k
        size_t run_est = std::min(num_rows, non_null_count * 2 + 1);
77
5.11k
        size_t run_bytes_est = run_est * kBytesPerRun + kReserveSlackBytes;
78
5.11k
        size_t raw_overhead = raw_bytes / 63 + 1;
79
5.11k
        size_t raw_est = raw_bytes + raw_overhead + kReserveSlackBytes;
80
5.11k
        size_t reserve_bytes = std::min(raw_est, run_bytes_est);
81
5.11k
        if (_bitmap_buf.capacity() < reserve_bytes) {
82
820
            const size_t cap = _bitmap_buf.capacity();
83
820
            const size_t grow = cap + cap / 2;
84
820
            const size_t new_cap = std::max(reserve_bytes, grow);
85
820
            _bitmap_buf.reserve(new_cap);
86
820
        }
87
5.11k
    }
88
89
802k
    void add_run(bool value, size_t run) {
90
802k
        _has_null |= value;
91
802k
        _rle_encoder.Put(value, run);
92
802k
    }
93
94
    // Returns whether the building nullmap contains nullptr
95
77.0k
    bool has_null() const { return _has_null; }
96
97
9.40k
    Status finish(OwnedSlice* slice) {
98
9.40k
        _rle_encoder.Flush();
99
9.40k
        RETURN_IF_CATCH_EXCEPTION({ *slice = _bitmap_buf.build(); });
100
9.40k
        return Status::OK();
101
9.40k
    }
102
103
77.0k
    void reset() {
104
77.0k
        _has_null = false;
105
77.0k
        _rle_encoder.Clear();
106
77.0k
    }
107
108
25.9k
    uint64_t size() { return _bitmap_buf.size(); }
109
110
private:
111
    static constexpr size_t kSmallReserveBytes = 64;
112
    static constexpr size_t kReserveSlackBytes = 16;
113
    static constexpr size_t kBytesPerRun = 6;
114
115
    bool _has_null;
116
    faststring _bitmap_buf;
117
    RleEncoder<bool> _rle_encoder;
118
};
119
120
inline ScalarColumnWriter* get_null_writer(const ColumnWriterOptions& opts,
121
3.96k
                                           io::FileWriter* file_writer, uint32_t id) {
122
3.96k
    if (!opts.meta->is_nullable()) {
123
509
        return nullptr;
124
509
    }
125
126
3.45k
    FieldType null_type = FieldType::OLAP_FIELD_TYPE_TINYINT;
127
3.45k
    ColumnWriterOptions null_options;
128
3.45k
    null_options.meta = opts.meta->add_children_columns();
129
3.45k
    null_options.meta->set_column_id(id);
130
3.45k
    null_options.meta->set_unique_id(id);
131
3.45k
    null_options.meta->set_type(int(null_type));
132
3.45k
    null_options.meta->set_is_nullable(false);
133
3.45k
    null_options.meta->set_length(
134
3.45k
            cast_set<int32_t>(field_type_size(FieldType::OLAP_FIELD_TYPE_TINYINT)));
135
3.45k
    null_options.meta->set_encoding(DEFAULT_ENCODING);
136
3.45k
    null_options.meta->set_compression(opts.meta->compression());
137
138
3.45k
    null_options.need_zone_map = false;
139
3.45k
    null_options.need_bloom_filter = false;
140
3.45k
    null_options.encoding_preference = opts.encoding_preference;
141
142
3.45k
    TabletColumn null_column =
143
3.45k
            TabletColumn(FieldAggregationMethod::OLAP_FIELD_AGGREGATION_NONE, null_type, false,
144
3.45k
                         null_options.meta->unique_id(), null_options.meta->length());
145
3.45k
    null_column.set_name("nullable");
146
3.45k
    null_column.set_index_length(-1); // no short key index
147
3.45k
    std::unique_ptr<StorageField> null_field(StorageFieldFactory::create(null_column));
148
3.45k
    return new ScalarColumnWriter(null_options, std::move(null_field), file_writer);
149
3.96k
}
150
151
ColumnWriter::ColumnWriter(std::unique_ptr<StorageField> field, bool is_nullable,
152
                           ColumnMetaPB* meta)
153
132k
        : _field(std::move(field)), _is_nullable(is_nullable), _column_meta(meta) {
154
132k
    _data_type = DataTypeFactory::instance().create_data_type(*_column_meta);
155
132k
}
156
Status ColumnWriter::create_struct_writer(const ColumnWriterOptions& opts,
157
                                          const TabletColumn* column, io::FileWriter* file_writer,
158
230
                                          std::unique_ptr<ColumnWriter>* writer) {
159
    // not support empty struct
160
230
    DCHECK(column->get_subtype_count() >= 1);
161
230
    std::vector<std::unique_ptr<ColumnWriter>> sub_column_writers;
162
230
    sub_column_writers.reserve(column->get_subtype_count());
163
1.61k
    for (uint32_t i = 0; i < column->get_subtype_count(); i++) {
164
1.38k
        const TabletColumn& sub_column = column->get_sub_column(i);
165
1.38k
        RETURN_IF_ERROR(sub_column.check_valid());
166
167
        // create sub writer
168
1.38k
        ColumnWriterOptions column_options;
169
1.38k
        column_options.meta = opts.meta->mutable_children_columns(i);
170
1.38k
        column_options.need_zone_map = false;
171
1.38k
        column_options.need_bloom_filter = sub_column.is_bf_column();
172
1.38k
        column_options.encoding_preference = opts.encoding_preference;
173
1.38k
        std::unique_ptr<ColumnWriter> sub_column_writer;
174
1.38k
        RETURN_IF_ERROR(
175
1.38k
                ColumnWriter::create(column_options, &sub_column, file_writer, &sub_column_writer));
176
1.38k
        sub_column_writers.push_back(std::move(sub_column_writer));
177
1.38k
    }
178
179
230
    ScalarColumnWriter* null_writer =
180
230
            get_null_writer(opts, file_writer, column->get_subtype_count() + 1);
181
182
230
    *writer = std::unique_ptr<ColumnWriter>(new StructColumnWriter(
183
230
            opts, std::unique_ptr<StorageField>(StorageFieldFactory::create(*column)), null_writer,
184
230
            sub_column_writers));
185
230
    return Status::OK();
186
230
}
187
188
Status ColumnWriter::create_array_writer(const ColumnWriterOptions& opts,
189
                                         const TabletColumn* column, io::FileWriter* file_writer,
190
2.06k
                                         std::unique_ptr<ColumnWriter>* writer) {
191
2.06k
    DCHECK(column->get_subtype_count() == 1);
192
2.06k
    const TabletColumn& item_column = column->get_sub_column(0);
193
2.06k
    RETURN_IF_ERROR(item_column.check_valid());
194
195
    // create item writer
196
2.06k
    ColumnWriterOptions item_options;
197
2.06k
    item_options.meta = opts.meta->mutable_children_columns(0);
198
2.06k
    item_options.need_zone_map = false;
199
2.06k
    item_options.need_bloom_filter = item_column.is_bf_column();
200
2.06k
    item_options.encoding_preference = opts.encoding_preference;
201
2.06k
    std::unique_ptr<ColumnWriter> item_writer;
202
2.06k
    RETURN_IF_ERROR(ColumnWriter::create(item_options, &item_column, file_writer, &item_writer));
203
204
    // create length writer
205
2.06k
    FieldType length_type = FieldType::OLAP_FIELD_TYPE_UNSIGNED_BIGINT;
206
207
2.06k
    ColumnWriterOptions length_options;
208
2.06k
    length_options.meta = opts.meta->add_children_columns();
209
2.06k
    length_options.meta->set_column_id(2);
210
2.06k
    length_options.meta->set_unique_id(2);
211
2.06k
    length_options.meta->set_type(int(length_type));
212
2.06k
    length_options.meta->set_is_nullable(false);
213
2.06k
    length_options.meta->set_length(
214
2.06k
            cast_set<int32_t>(field_type_size(FieldType::OLAP_FIELD_TYPE_UNSIGNED_BIGINT)));
215
2.06k
    length_options.meta->set_encoding(DEFAULT_ENCODING);
216
2.06k
    length_options.meta->set_compression(opts.meta->compression());
217
218
2.06k
    length_options.need_zone_map = false;
219
2.06k
    length_options.need_bloom_filter = false;
220
2.06k
    length_options.encoding_preference = opts.encoding_preference;
221
222
2.06k
    TabletColumn length_column =
223
2.06k
            TabletColumn(FieldAggregationMethod::OLAP_FIELD_AGGREGATION_NONE, length_type,
224
2.06k
                         length_options.meta->is_nullable(), length_options.meta->unique_id(),
225
2.06k
                         length_options.meta->length());
226
2.06k
    length_column.set_name("length");
227
2.06k
    length_column.set_index_length(-1); // no short key index
228
2.06k
    std::unique_ptr<StorageField> bigint_field(StorageFieldFactory::create(length_column));
229
2.06k
    auto* length_writer =
230
2.06k
            new OffsetColumnWriter(length_options, std::move(bigint_field), file_writer);
231
232
2.06k
    ScalarColumnWriter* null_writer = get_null_writer(opts, file_writer, 3);
233
234
2.06k
    *writer = std::unique_ptr<ColumnWriter>(new ArrayColumnWriter(
235
2.06k
            opts, std::unique_ptr<StorageField>(StorageFieldFactory::create(*column)),
236
2.06k
            length_writer, null_writer, std::move(item_writer)));
237
2.06k
    return Status::OK();
238
2.06k
}
239
240
Status ColumnWriter::create_map_writer(const ColumnWriterOptions& opts, const TabletColumn* column,
241
                                       io::FileWriter* file_writer,
242
1.67k
                                       std::unique_ptr<ColumnWriter>* writer) {
243
1.67k
    DCHECK(column->get_subtype_count() == 2);
244
1.67k
    if (column->get_subtype_count() < 2) {
245
0
        return Status::InternalError(
246
0
                "If you upgraded from version 1.2.*, please DROP the MAP columns and then "
247
0
                "ADD the MAP columns back.");
248
0
    }
249
    // create key & value writer
250
1.67k
    std::vector<std::unique_ptr<ColumnWriter>> inner_writer_list;
251
5.02k
    for (int i = 0; i < 2; ++i) {
252
3.34k
        const TabletColumn& item_column = column->get_sub_column(i);
253
3.34k
        RETURN_IF_ERROR(item_column.check_valid());
254
255
        // create item writer
256
3.34k
        ColumnWriterOptions item_options;
257
3.34k
        item_options.meta = opts.meta->mutable_children_columns(i);
258
3.34k
        item_options.need_zone_map = false;
259
3.34k
        item_options.need_bloom_filter = item_column.is_bf_column();
260
3.34k
        item_options.encoding_preference = opts.encoding_preference;
261
3.34k
        std::unique_ptr<ColumnWriter> item_writer;
262
3.34k
        RETURN_IF_ERROR(
263
3.34k
                ColumnWriter::create(item_options, &item_column, file_writer, &item_writer));
264
3.34k
        inner_writer_list.push_back(std::move(item_writer));
265
3.34k
    }
266
267
    // create offset writer
268
1.67k
    FieldType length_type = FieldType::OLAP_FIELD_TYPE_UNSIGNED_BIGINT;
269
270
    // Be Cautious: column unique id is used for column reader creation
271
1.67k
    ColumnWriterOptions length_options;
272
1.67k
    length_options.meta = opts.meta->add_children_columns();
273
1.67k
    length_options.meta->set_column_id(column->get_subtype_count() + 1);
274
1.67k
    length_options.meta->set_unique_id(column->get_subtype_count() + 1);
275
1.67k
    length_options.meta->set_type(int(length_type));
276
1.67k
    length_options.meta->set_is_nullable(false);
277
1.67k
    length_options.meta->set_length(
278
1.67k
            cast_set<int32_t>(field_type_size(FieldType::OLAP_FIELD_TYPE_UNSIGNED_BIGINT)));
279
1.67k
    length_options.meta->set_encoding(DEFAULT_ENCODING);
280
1.67k
    length_options.meta->set_compression(opts.meta->compression());
281
282
1.67k
    length_options.need_zone_map = false;
283
1.67k
    length_options.need_bloom_filter = false;
284
1.67k
    length_options.encoding_preference = opts.encoding_preference;
285
286
1.67k
    TabletColumn length_column =
287
1.67k
            TabletColumn(FieldAggregationMethod::OLAP_FIELD_AGGREGATION_NONE, length_type,
288
1.67k
                         length_options.meta->is_nullable(), length_options.meta->unique_id(),
289
1.67k
                         length_options.meta->length());
290
1.67k
    length_column.set_name("length");
291
1.67k
    length_column.set_index_length(-1); // no short key index
292
1.67k
    std::unique_ptr<StorageField> bigint_field(StorageFieldFactory::create(length_column));
293
1.67k
    auto* length_writer =
294
1.67k
            new OffsetColumnWriter(length_options, std::move(bigint_field), file_writer);
295
296
1.67k
    ScalarColumnWriter* null_writer =
297
1.67k
            get_null_writer(opts, file_writer, column->get_subtype_count() + 2);
298
299
1.67k
    *writer = std::unique_ptr<ColumnWriter>(new MapColumnWriter(
300
1.67k
            opts, std::unique_ptr<StorageField>(StorageFieldFactory::create(*column)), null_writer,
301
1.67k
            length_writer, inner_writer_list));
302
303
1.67k
    return Status::OK();
304
1.67k
}
305
306
Status ColumnWriter::create_agg_state_writer(const ColumnWriterOptions& opts,
307
                                             const TabletColumn* column,
308
                                             io::FileWriter* file_writer,
309
8
                                             std::unique_ptr<ColumnWriter>* writer) {
310
8
    auto data_type = DataTypeFactory::instance().create_data_type(*column);
311
8
    const auto* agg_state_type = assert_cast<const DataTypeAggState*>(data_type.get());
312
8
    auto type = agg_state_type->get_serialized_type()->get_primitive_type();
313
8
    if (type == PrimitiveType::TYPE_STRING || type == PrimitiveType::INVALID_TYPE ||
314
8
        type == PrimitiveType::TYPE_FIXED_LENGTH_OBJECT || type == PrimitiveType::TYPE_BITMAP) {
315
8
        *writer = std::unique_ptr<ColumnWriter>(new ScalarColumnWriter(
316
8
                opts, std::unique_ptr<StorageField>(StorageFieldFactory::create(*column)),
317
8
                file_writer));
318
8
    } else if (type == PrimitiveType::TYPE_ARRAY) {
319
0
        RETURN_IF_ERROR(create_array_writer(opts, column, file_writer, writer));
320
0
    } else if (type == PrimitiveType::TYPE_MAP) {
321
0
        RETURN_IF_ERROR(create_map_writer(opts, column, file_writer, writer));
322
0
    } else {
323
0
        throw Exception(ErrorCode::INTERNAL_ERROR,
324
0
                        "OLAP_FIELD_TYPE_AGG_STATE meet unsupported type: {}",
325
0
                        agg_state_type->get_name());
326
0
    }
327
8
    return Status::OK();
328
8
}
329
330
Status ColumnWriter::create_variant_writer(const ColumnWriterOptions& opts,
331
                                           const TabletColumn* column, io::FileWriter* file_writer,
332
319
                                           std::unique_ptr<ColumnWriter>* writer) {
333
    // Variant extracted columns have two kinds of physical writers:
334
    // - Doc-value snapshot column (`...__DORIS_VARIANT_DOC_VALUE__...`): use `VariantDocCompactWriter`
335
    //   to store the doc snapshot in a compact binary form.
336
    // - Regular extracted subcolumns: use `VariantSubcolumnWriter`.
337
    // The root VARIANT column itself uses `VariantColumnWriter`.
338
319
    if (column->is_extracted_column()) {
339
3
        if (column->name().find(DOC_VALUE_COLUMN_PATH) != std::string::npos) {
340
2
            *writer = std::make_unique<VariantDocCompactWriter>(
341
2
                    opts, column,
342
2
                    std::unique_ptr<StorageField>(StorageFieldFactory::create(*column)));
343
2
            return Status::OK();
344
2
        }
345
1
        VLOG_DEBUG << "gen subwriter for " << column->path_info_ptr()->get_path();
346
1
        *writer = std::make_unique<VariantSubcolumnWriter>(
347
1
                opts, column, std::unique_ptr<StorageField>(StorageFieldFactory::create(*column)));
348
1
        return Status::OK();
349
3
    }
350
316
    *writer = std::make_unique<VariantColumnWriter>(
351
316
            opts, column, std::unique_ptr<StorageField>(StorageFieldFactory::create(*column)));
352
316
    return Status::OK();
353
319
}
354
355
//Todo(Amory): here should according nullable and offset and need sub to simply this function
356
Status ColumnWriter::create(const ColumnWriterOptions& opts, const TabletColumn* column,
357
124k
                            io::FileWriter* file_writer, std::unique_ptr<ColumnWriter>* writer) {
358
124k
    std::unique_ptr<StorageField> field(StorageFieldFactory::create(*column));
359
124k
    DCHECK(field.get() != nullptr);
360
124k
    if (is_scalar_type(column->type())) {
361
120k
        *writer = std::unique_ptr<ColumnWriter>(
362
120k
                new ScalarColumnWriter(opts, std::move(field), file_writer));
363
120k
        return Status::OK();
364
120k
    } else {
365
3.94k
        switch (column->type()) {
366
8
        case FieldType::OLAP_FIELD_TYPE_AGG_STATE: {
367
8
            RETURN_IF_ERROR(create_agg_state_writer(opts, column, file_writer, writer));
368
8
            return Status::OK();
369
8
        }
370
230
        case FieldType::OLAP_FIELD_TYPE_STRUCT: {
371
230
            RETURN_IF_ERROR(create_struct_writer(opts, column, file_writer, writer));
372
230
            return Status::OK();
373
230
        }
374
2.06k
        case FieldType::OLAP_FIELD_TYPE_ARRAY: {
375
2.06k
            RETURN_IF_ERROR(create_array_writer(opts, column, file_writer, writer));
376
2.06k
            return Status::OK();
377
2.06k
        }
378
1.34k
        case FieldType::OLAP_FIELD_TYPE_MAP: {
379
1.34k
            RETURN_IF_ERROR(create_map_writer(opts, column, file_writer, writer));
380
1.34k
            return Status::OK();
381
1.34k
        }
382
318
        case FieldType::OLAP_FIELD_TYPE_VARIANT: {
383
            // Process columns with sparse column
384
318
            RETURN_IF_ERROR(create_variant_writer(opts, column, file_writer, writer));
385
318
            return Status::OK();
386
318
        }
387
0
        default:
388
0
            return Status::NotSupported("unsupported type for ColumnWriter: {}",
389
0
                                        std::to_string(int(field->type())));
390
3.94k
        }
391
3.94k
    }
392
124k
}
393
394
Status ColumnWriter::append_nullable(const uint8_t* is_null_bits, const void* data,
395
651k
                                     size_t num_rows) {
396
651k
    const auto* ptr = (const uint8_t*)data;
397
651k
    BitmapIterator null_iter(is_null_bits, num_rows);
398
651k
    bool is_null = false;
399
651k
    size_t this_run = 0;
400
1.30M
    while ((this_run = null_iter.Next(&is_null)) > 0) {
401
651k
        if (is_null) {
402
0
            RETURN_IF_ERROR(append_nulls(this_run));
403
651k
        } else {
404
651k
            RETURN_IF_ERROR(append_data(&ptr, this_run));
405
651k
        }
406
651k
    }
407
651k
    return Status::OK();
408
651k
}
409
410
Status ColumnWriter::append_nullable(const uint8_t* null_map, const uint8_t** ptr,
411
0
                                     size_t num_rows) {
412
    // Fast path: use SIMD to detect all-NULL or all-non-NULL columns
413
0
    if (config::enable_rle_batch_put_optimization) {
414
0
        size_t non_null_count =
415
0
                simd::count_zero_num(reinterpret_cast<const int8_t*>(null_map), num_rows);
416
417
0
        if (non_null_count == 0) {
418
            // All NULL: skip run-length iteration, directly append all nulls
419
0
            RETURN_IF_ERROR(append_nulls(num_rows));
420
0
            *ptr += get_field()->size() * num_rows;
421
0
            return Status::OK();
422
0
        }
423
424
0
        if (non_null_count == num_rows) {
425
            // All non-NULL: skip run-length iteration, directly append all data
426
0
            return append_data(ptr, num_rows);
427
0
        }
428
0
    }
429
430
    // Mixed case or sparse optimization disabled: use run-length processing
431
0
    size_t offset = 0;
432
0
    auto next_run_step = [&]() {
433
0
        size_t step = 1;
434
0
        for (auto i = offset + 1; i < num_rows; ++i) {
435
0
            if (null_map[offset] == null_map[i]) {
436
0
                step++;
437
0
            } else {
438
0
                break;
439
0
            }
440
0
        }
441
0
        return step;
442
0
    };
443
444
0
    do {
445
0
        auto step = next_run_step();
446
0
        if (null_map[offset]) {
447
0
            RETURN_IF_ERROR(append_nulls(step));
448
0
            *ptr += get_field()->size() * step;
449
0
        } else {
450
            // TODO:
451
            //  1. `*ptr += get_field()->size() * step;` should do in this function, not append_data;
452
            //  2. support array vectorized load and ptr offset add
453
0
            RETURN_IF_ERROR(append_data(ptr, step));
454
0
        }
455
0
        offset += step;
456
0
    } while (offset < num_rows);
457
458
0
    return Status::OK();
459
0
}
460
461
126k
Status ColumnWriter::append(const uint8_t* nullmap, const void* data, size_t num_rows) {
462
126k
    assert(data && num_rows > 0);
463
126k
    const auto* ptr = (const uint8_t*)data;
464
126k
    if (nullmap) {
465
76.6k
        return append_nullable(nullmap, &ptr, num_rows);
466
76.6k
    } else {
467
50.0k
        return append_data(&ptr, num_rows);
468
50.0k
    }
469
126k
}
470
471
///////////////////////////////////////////////////////////////////////////////////
472
473
ScalarColumnWriter::ScalarColumnWriter(const ColumnWriterOptions& opts,
474
                                       std::unique_ptr<StorageField> field,
475
                                       io::FileWriter* file_writer)
476
127k
        : ColumnWriter(std::move(field), opts.meta->is_nullable(), opts.meta),
477
127k
          _opts(opts),
478
127k
          _file_writer(file_writer),
479
127k
          _data_size(0) {
480
    // these opts.meta fields should be set by client
481
127k
    DCHECK(opts.meta->has_column_id());
482
127k
    DCHECK(opts.meta->has_unique_id());
483
127k
    DCHECK(opts.meta->has_type());
484
127k
    DCHECK(opts.meta->has_length());
485
127k
    DCHECK(opts.meta->has_encoding());
486
127k
    DCHECK(opts.meta->has_compression());
487
127k
    DCHECK(opts.meta->has_is_nullable());
488
127k
    DCHECK(file_writer != nullptr);
489
127k
    _inverted_index_builders.resize(_opts.inverted_indexes.size());
490
127k
}
491
492
128k
ScalarColumnWriter::~ScalarColumnWriter() {
493
    // delete all pages
494
128k
    _pages.clear();
495
128k
}
496
497
127k
Status ScalarColumnWriter::init() {
498
127k
    RETURN_IF_ERROR(get_block_compression_codec(_opts.meta->compression(), &_compress_codec));
499
500
127k
    PageBuilder* page_builder = nullptr;
501
502
127k
    RETURN_IF_ERROR(EncodingInfo::get(get_field()->type(), _opts.meta->encoding(),
503
127k
                                      _opts.encoding_preference, &_encoding_info));
504
127k
    _opts.meta->set_encoding(_encoding_info->encoding());
505
    // create page builder
506
127k
    PageBuilderOptions opts;
507
127k
    opts.data_page_size = _opts.data_page_size;
508
127k
    opts.dict_page_size = _opts.dict_page_size;
509
127k
    opts.encoding_preference = _opts.encoding_preference;
510
127k
    RETURN_IF_ERROR(_encoding_info->create_page_builder(opts, &page_builder));
511
127k
    if (page_builder == nullptr) {
512
0
        return Status::NotSupported("Failed to create page builder for type {} and encoding {}",
513
0
                                    get_field()->type(), _opts.meta->encoding());
514
0
    }
515
    // should store more concrete encoding type instead of DEFAULT_ENCODING
516
    // because the default encoding of a data type can be changed in the future
517
127k
    DCHECK_NE(_opts.meta->encoding(), DEFAULT_ENCODING);
518
18.4E
    VLOG_DEBUG << fmt::format(
519
18.4E
            "[verbose] scalar column writer init, column_id={}, type={}, encoding={}, "
520
18.4E
            "is_nullable={}",
521
18.4E
            _opts.meta->column_id(), get_field()->type(),
522
18.4E
            EncodingTypePB_Name(_opts.meta->encoding()), _opts.meta->is_nullable());
523
127k
    _page_builder.reset(page_builder);
524
    // create ordinal builder
525
127k
    _ordinal_index_builder = std::make_unique<OrdinalIndexWriter>();
526
    // create null bitmap builder
527
127k
    if (is_nullable()) {
528
73.0k
        _null_bitmap_builder = std::make_unique<NullBitmapBuilder>();
529
73.0k
    }
530
127k
    if (_opts.need_zone_map) {
531
112k
        RETURN_IF_ERROR(
532
112k
                ZoneMapIndexWriter::create(_data_type, get_field(), _zone_map_index_builder));
533
112k
    }
534
535
127k
    if (_opts.need_inverted_index) {
536
2.18k
        do {
537
4.38k
            for (size_t i = 0; i < _opts.inverted_indexes.size(); i++) {
538
2.19k
                DBUG_EXECUTE_IF("column_writer.init", {
539
2.19k
                    class InvertedIndexColumnWriterEmpty final : public IndexColumnWriter {
540
2.19k
                    public:
541
2.19k
                        Status init() override { return Status::OK(); }
542
2.19k
                        Status add_values(const std::string name, const void* values,
543
2.19k
                                          size_t count) override {
544
2.19k
                            return Status::OK();
545
2.19k
                        }
546
2.19k
                        Status add_array_values(size_t field_size, const CollectionValue* values,
547
2.19k
                                                size_t count) override {
548
2.19k
                            return Status::OK();
549
2.19k
                        }
550
2.19k
                        Status add_array_values(size_t field_size, const void* value_ptr,
551
2.19k
                                                const uint8_t* null_map, const uint8_t* offsets_ptr,
552
2.19k
                                                size_t count) override {
553
2.19k
                            return Status::OK();
554
2.19k
                        }
555
2.19k
                        Status add_nulls(uint32_t count) override { return Status::OK(); }
556
2.19k
                        Status add_array_nulls(const uint8_t* null_map, size_t num_rows) override {
557
2.19k
                            return Status::OK();
558
2.19k
                        }
559
2.19k
                        Status finish() override { return Status::OK(); }
560
2.19k
                        int64_t size() const override { return 0; }
561
2.19k
                        void close_on_error() override {}
562
2.19k
                    };
563
564
2.19k
                    _inverted_index_builders[i] =
565
2.19k
                            std::make_unique<InvertedIndexColumnWriterEmpty>();
566
567
2.19k
                    break;
568
2.19k
                });
569
570
2.19k
                RETURN_IF_ERROR(IndexColumnWriter::create(get_field(), &_inverted_index_builders[i],
571
2.19k
                                                          _opts.index_file_writer,
572
2.19k
                                                          _opts.inverted_indexes[i]));
573
2.19k
            }
574
2.18k
        } while (false);
575
2.18k
    }
576
127k
    if (_opts.need_bloom_filter) {
577
7
        if (_opts.is_ngram_bf_index) {
578
0
            RETURN_IF_ERROR(NGramBloomFilterIndexWriterImpl::create(
579
0
                    BloomFilterOptions(), get_field()->type(), _opts.gram_size, _opts.gram_bf_size,
580
0
                    &_bloom_filter_index_builder));
581
7
        } else {
582
7
            RETURN_IF_ERROR(BloomFilterIndexWriter::create(_opts.bf_options, get_field()->type(),
583
7
                                                           &_bloom_filter_index_builder));
584
7
        }
585
7
    }
586
127k
    return Status::OK();
587
127k
}
588
589
43.9k
Status ScalarColumnWriter::append_nulls(size_t num_rows) {
590
43.9k
    _null_bitmap_builder->add_run(true, num_rows);
591
43.9k
    _next_rowid += num_rows;
592
43.9k
    if (_opts.need_zone_map) {
593
36.3k
        _zone_map_index_builder->add_nulls(cast_set<uint32_t>(num_rows));
594
36.3k
    }
595
43.9k
    if (_opts.need_inverted_index) {
596
2
        for (const auto& builder : _inverted_index_builders) {
597
2
            RETURN_IF_ERROR(builder->add_nulls(cast_set<uint32_t>(num_rows)));
598
2
        }
599
2
    }
600
43.9k
    if (_opts.need_bloom_filter) {
601
0
        _bloom_filter_index_builder->add_nulls(cast_set<uint32_t>(num_rows));
602
0
    }
603
43.9k
    return Status::OK();
604
43.9k
}
605
606
// append data to page builder. this function will make sure that
607
// num_rows must be written before return. And ptr will be modified
608
// to next data should be written
609
806k
Status ScalarColumnWriter::append_data(const uint8_t** ptr, size_t num_rows) {
610
806k
    size_t remaining = num_rows;
611
1.61M
    while (remaining > 0) {
612
812k
        size_t num_written = remaining;
613
812k
        RETURN_IF_ERROR(append_data_in_current_page(ptr, &num_written));
614
615
812k
        remaining -= num_written;
616
617
812k
        if (_page_builder->is_page_full()) {
618
5.54k
            RETURN_IF_ERROR(finish_current_page());
619
5.54k
        }
620
812k
    }
621
806k
    return Status::OK();
622
806k
}
623
624
Status ScalarColumnWriter::_internal_append_data_in_current_page(const uint8_t* data,
625
816k
                                                                 size_t* num_written) {
626
816k
    RETURN_IF_ERROR(_page_builder->add(data, num_written));
627
816k
    if (_opts.need_zone_map) {
628
793k
        _zone_map_index_builder->add_values(data, *num_written);
629
793k
    }
630
816k
    if (_opts.need_inverted_index) {
631
4.48k
        for (const auto& builder : _inverted_index_builders) {
632
4.48k
            RETURN_IF_ERROR(builder->add_values(get_field()->name(), data, *num_written));
633
4.48k
        }
634
4.47k
    }
635
816k
    if (_opts.need_bloom_filter) {
636
7
        RETURN_IF_ERROR(_bloom_filter_index_builder->add_values(data, *num_written));
637
7
    }
638
639
816k
    _next_rowid += *num_written;
640
641
    // we must write null bits after write data, because we don't
642
    // know how many rows can be written into current page
643
816k
    if (is_nullable()) {
644
758k
        _null_bitmap_builder->add_run(false, *num_written);
645
758k
    }
646
816k
    return Status::OK();
647
816k
}
648
649
816k
Status ScalarColumnWriter::append_data_in_current_page(const uint8_t** data, size_t* num_written) {
650
816k
    RETURN_IF_ERROR(append_data_in_current_page(*data, num_written));
651
816k
    *data += get_field()->size() * (*num_written);
652
816k
    return Status::OK();
653
816k
}
654
655
Status ScalarColumnWriter::append_nullable(const uint8_t* null_map, const uint8_t** ptr,
656
72.9k
                                           size_t num_rows) {
657
    // When optimization is disabled, use base class implementation
658
72.9k
    if (!config::enable_rle_batch_put_optimization) {
659
0
        return ColumnWriter::append_nullable(null_map, ptr, num_rows);
660
0
    }
661
662
72.9k
    if (UNLIKELY(num_rows == 0)) {
663
0
        return Status::OK();
664
0
    }
665
666
    // Build run-length encoded null runs using memchr for fast boundary detection
667
72.9k
    _null_run_buffer.clear();
668
72.9k
    if (_null_run_buffer.capacity() < num_rows) {
669
72.0k
        _null_run_buffer.reserve(std::min(num_rows, size_t(256)));
670
72.0k
    }
671
672
72.9k
    size_t non_null_count = 0;
673
72.9k
    size_t offset = 0;
674
218k
    while (offset < num_rows) {
675
145k
        bool is_null = null_map[offset] != 0;
676
145k
        size_t remaining = num_rows - offset;
677
145k
        const uint8_t* run_end =
678
145k
                static_cast<const uint8_t*>(memchr(null_map + offset, is_null ? 0 : 1, remaining));
679
145k
        size_t run_length = run_end != nullptr ? (run_end - (null_map + offset)) : remaining;
680
145k
        _null_run_buffer.push_back(NullRun {is_null, static_cast<uint32_t>(run_length)});
681
145k
        if (!is_null) {
682
102k
            non_null_count += run_length;
683
102k
        }
684
145k
        offset += run_length;
685
145k
    }
686
687
    // Pre-allocate buffer based on estimated size
688
72.9k
    if (_null_bitmap_builder != nullptr) {
689
72.9k
        size_t current_rows = _next_rowid - _first_rowid;
690
72.9k
        size_t expected_rows = current_rows + num_rows;
691
72.9k
        size_t est_non_null = non_null_count;
692
72.9k
        if (num_rows > 0 && expected_rows > num_rows) {
693
2.98k
            est_non_null = (non_null_count * expected_rows) / num_rows;
694
2.98k
        }
695
72.9k
        _null_bitmap_builder->reserve_for_write(expected_rows, est_non_null);
696
72.9k
    }
697
698
72.9k
    if (non_null_count == 0) {
699
        // All NULL: skip data writing, only update null bitmap and indexes
700
5.80k
        RETURN_IF_ERROR(append_nulls(num_rows));
701
5.80k
        *ptr += get_field()->size() * num_rows;
702
5.80k
        return Status::OK();
703
5.80k
    }
704
705
67.1k
    if (non_null_count == num_rows) {
706
        // All non-NULL: use normal append_data which handles both data and null bitmap
707
63.5k
        return append_data(ptr, num_rows);
708
63.5k
    }
709
710
    // Process by runs
711
76.0k
    for (const auto& run : _null_run_buffer) {
712
76.0k
        size_t run_length = run.len;
713
76.0k
        if (run.is_null) {
714
37.2k
            RETURN_IF_ERROR(append_nulls(run_length));
715
37.2k
            *ptr += get_field()->size() * run_length;
716
38.8k
        } else {
717
            // TODO:
718
            //  1. `*ptr += get_field()->size() * step;` should do in this function, not append_data;
719
            //  2. support array vectorized load and ptr offset add
720
38.8k
            RETURN_IF_ERROR(append_data(ptr, run_length));
721
38.8k
        }
722
76.0k
    }
723
724
3.63k
    return Status::OK();
725
3.63k
}
726
727
49.5k
uint64_t ScalarColumnWriter::estimate_buffer_size() {
728
49.5k
    uint64_t size = _data_size;
729
49.5k
    size += _page_builder->size();
730
49.5k
    if (is_nullable()) {
731
25.9k
        size += _null_bitmap_builder->size();
732
25.9k
    }
733
49.5k
    size += _ordinal_index_builder->size();
734
49.5k
    if (_opts.need_zone_map) {
735
43.4k
        size += _zone_map_index_builder->size();
736
43.4k
    }
737
49.5k
    if (_opts.need_bloom_filter) {
738
8
        size += _bloom_filter_index_builder->size();
739
8
    }
740
49.5k
    return size;
741
49.5k
}
742
743
128k
Status ScalarColumnWriter::finish() {
744
128k
    RETURN_IF_ERROR(finish_current_page());
745
128k
    _opts.meta->set_num_rows(_next_rowid);
746
128k
    return Status::OK();
747
128k
}
748
749
128k
Status ScalarColumnWriter::write_data() {
750
128k
    auto offset = _file_writer->bytes_appended();
751
187k
    auto collect_uncompressed_bytes = [](const PageFooterPB& footer) {
752
187k
        return footer.uncompressed_size() + footer.ByteSizeLong() +
753
187k
               sizeof(uint32_t) /* footer size */ + sizeof(uint32_t) /* checksum */;
754
187k
    };
755
132k
    for (auto& page : _pages) {
756
132k
        _total_uncompressed_data_pages_size += collect_uncompressed_bytes(page->footer);
757
132k
        RETURN_IF_ERROR(_write_data_page(page.get()));
758
132k
    }
759
128k
    _pages.clear();
760
    // write column dict
761
128k
    if (_encoding_info->encoding() == DICT_ENCODING) {
762
55.1k
        OwnedSlice dict_body;
763
55.1k
        RETURN_IF_ERROR(_page_builder->get_dictionary_page(&dict_body));
764
55.1k
        EncodingTypePB dict_word_page_encoding;
765
55.1k
        RETURN_IF_ERROR(_page_builder->get_dictionary_page_encoding(&dict_word_page_encoding));
766
767
55.1k
        PageFooterPB footer;
768
55.1k
        footer.set_type(DICTIONARY_PAGE);
769
55.1k
        footer.set_uncompressed_size(cast_set<uint32_t>(dict_body.slice().get_size()));
770
55.1k
        footer.mutable_dict_page_footer()->set_encoding(dict_word_page_encoding);
771
55.1k
        _total_uncompressed_data_pages_size += collect_uncompressed_bytes(footer);
772
773
55.1k
        PagePointer dict_pp;
774
55.1k
        RETURN_IF_ERROR(PageIO::compress_and_write_page(
775
55.1k
                _compress_codec, _opts.compression_min_space_saving, _file_writer,
776
55.1k
                {dict_body.slice()}, footer, &dict_pp));
777
55.1k
        dict_pp.to_proto(_opts.meta->mutable_dict_page());
778
55.1k
    }
779
128k
    _total_compressed_data_pages_size += _file_writer->bytes_appended() - offset;
780
128k
    _page_builder.reset();
781
128k
    return Status::OK();
782
128k
}
783
784
127k
Status ScalarColumnWriter::write_ordinal_index() {
785
127k
    return _ordinal_index_builder->finish(_file_writer, _opts.meta->add_indexes());
786
127k
}
787
788
113k
Status ScalarColumnWriter::write_zone_map() {
789
113k
    if (_opts.need_zone_map) {
790
112k
        return _zone_map_index_builder->finish(_file_writer, _opts.meta->add_indexes());
791
112k
    }
792
1.14k
    return Status::OK();
793
113k
}
794
795
112k
Status ScalarColumnWriter::write_inverted_index() {
796
112k
    if (_opts.need_inverted_index) {
797
2.19k
        for (const auto& builder : _inverted_index_builders) {
798
2.19k
            RETURN_IF_ERROR(builder->finish());
799
2.19k
        }
800
2.18k
    }
801
112k
    return Status::OK();
802
112k
}
803
804
112k
Status ScalarColumnWriter::write_bloom_filter_index() {
805
112k
    if (_opts.need_bloom_filter) {
806
7
        return _bloom_filter_index_builder->finish(_file_writer, _opts.meta->add_indexes());
807
7
    }
808
112k
    return Status::OK();
809
112k
}
810
811
// write a data page into file and update ordinal index
812
132k
Status ScalarColumnWriter::_write_data_page(Page* page) {
813
132k
    PagePointer pp;
814
132k
    std::vector<Slice> compressed_body;
815
250k
    for (auto& data : page->data) {
816
250k
        compressed_body.push_back(data.slice());
817
250k
    }
818
132k
    RETURN_IF_ERROR(PageIO::write_page(_file_writer, compressed_body, page->footer, &pp));
819
132k
    _ordinal_index_builder->append_entry(page->footer.data_page_footer().first_ordinal(), pp);
820
132k
    return Status::OK();
821
132k
}
822
823
133k
Status ScalarColumnWriter::finish_current_page() {
824
133k
    if (_next_rowid == _first_rowid) {
825
994
        return Status::OK();
826
994
    }
827
132k
    if (_opts.need_zone_map) {
828
        // If the number of rows in the current page is less than the threshold,
829
        // we will invalidate zone map index for this page by set pass_all to true.
830
117k
        if (_next_rowid - _first_rowid < config::zone_map_row_num_threshold) {
831
81.9k
            _zone_map_index_builder->invalid_page_zone_map();
832
81.9k
        }
833
117k
        RETURN_IF_ERROR(_zone_map_index_builder->flush());
834
117k
    }
835
836
132k
    if (_opts.need_bloom_filter) {
837
7
        RETURN_IF_ERROR(_bloom_filter_index_builder->flush());
838
7
    }
839
840
132k
    _raw_data_bytes += _page_builder->get_raw_data_size();
841
842
    // build data page body : encoded values + [nullmap]
843
132k
    std::vector<Slice> body;
844
132k
    OwnedSlice encoded_values;
845
132k
    RETURN_IF_ERROR(_page_builder->finish(&encoded_values));
846
132k
    RETURN_IF_ERROR(_page_builder->reset());
847
132k
    body.push_back(encoded_values.slice());
848
849
132k
    OwnedSlice nullmap;
850
132k
    if (_null_bitmap_builder != nullptr) {
851
77.0k
        if (is_nullable() && _null_bitmap_builder->has_null()) {
852
9.40k
            RETURN_IF_ERROR(_null_bitmap_builder->finish(&nullmap));
853
9.40k
            body.push_back(nullmap.slice());
854
9.40k
        }
855
77.0k
        _null_bitmap_builder->reset();
856
77.0k
    }
857
858
    // prepare data page footer
859
132k
    std::unique_ptr<Page> page(new Page());
860
132k
    page->footer.set_type(DATA_PAGE);
861
132k
    page->footer.set_uncompressed_size(cast_set<uint32_t>(Slice::compute_total_size(body)));
862
132k
    auto* data_page_footer = page->footer.mutable_data_page_footer();
863
132k
    data_page_footer->set_first_ordinal(_first_rowid);
864
132k
    data_page_footer->set_num_values(_next_rowid - _first_rowid);
865
132k
    data_page_footer->set_nullmap_size(cast_set<uint32_t>(nullmap.slice().size));
866
132k
    if (_new_page_callback != nullptr) {
867
3.77k
        _new_page_callback->put_extra_info_in_page(data_page_footer);
868
3.77k
    }
869
    // trying to compress page body
870
132k
    OwnedSlice compressed_body;
871
132k
    RETURN_IF_ERROR(PageIO::compress_page_body(_compress_codec, _opts.compression_min_space_saving,
872
132k
                                               body, &compressed_body));
873
132k
    if (compressed_body.slice().empty()) {
874
        // page body is uncompressed
875
117k
        page->data.emplace_back(std::move(encoded_values));
876
117k
        page->data.emplace_back(std::move(nullmap));
877
117k
    } else {
878
        // page body is compressed
879
14.6k
        page->data.emplace_back(std::move(compressed_body));
880
14.6k
    }
881
882
132k
    _push_back_page(std::move(page));
883
132k
    _first_rowid = _next_rowid;
884
132k
    return Status::OK();
885
132k
}
886
887
////////////////////////////////////////////////////////////////////////////////
888
889
////////////////////////////////////////////////////////////////////////////////
890
// offset column writer
891
////////////////////////////////////////////////////////////////////////////////
892
893
OffsetColumnWriter::OffsetColumnWriter(const ColumnWriterOptions& opts,
894
                                       std::unique_ptr<StorageField> field,
895
                                       io::FileWriter* file_writer)
896
3.73k
        : ScalarColumnWriter(opts, std::move(field), file_writer) {
897
    // now we only explain data in offset column as uint64
898
3.73k
    DCHECK(get_field()->type() == FieldType::OLAP_FIELD_TYPE_UNSIGNED_BIGINT);
899
3.73k
}
900
901
3.73k
OffsetColumnWriter::~OffsetColumnWriter() = default;
902
903
3.73k
Status OffsetColumnWriter::init() {
904
3.73k
    RETURN_IF_ERROR(ScalarColumnWriter::init());
905
3.73k
    register_flush_page_callback(this);
906
3.73k
    _next_offset = 0;
907
3.73k
    return Status::OK();
908
3.73k
}
909
910
3.97k
Status OffsetColumnWriter::append_data(const uint8_t** ptr, size_t num_rows) {
911
3.97k
    size_t remaining = num_rows;
912
8.01k
    while (remaining > 0) {
913
4.04k
        size_t num_written = remaining;
914
4.04k
        RETURN_IF_ERROR(append_data_in_current_page(ptr, &num_written));
915
        // Callers provide one extra tail offset after the written rows so the page footer can
916
        // store the next array item ordinal for the current page.
917
4.04k
        _next_offset = *(const uint64_t*)(*ptr);
918
4.04k
        remaining -= num_written;
919
920
4.04k
        if (_page_builder->is_page_full()) {
921
            // get next data for next array_item_rowid
922
72
            RETURN_IF_ERROR(finish_current_page());
923
72
        }
924
4.04k
    }
925
3.97k
    return Status::OK();
926
3.97k
}
927
928
3.77k
void OffsetColumnWriter::put_extra_info_in_page(DataPageFooterPB* footer) {
929
3.77k
    footer->set_next_array_item_ordinal(_next_offset);
930
3.77k
}
931
932
StructColumnWriter::StructColumnWriter(
933
        const ColumnWriterOptions& opts, std::unique_ptr<StorageField> field,
934
        ScalarColumnWriter* null_writer,
935
        std::vector<std::unique_ptr<ColumnWriter>>& sub_column_writers)
936
230
        : ColumnWriter(std::move(field), opts.meta->is_nullable(), opts.meta), _opts(opts) {
937
1.38k
    for (auto& sub_column_writer : sub_column_writers) {
938
1.38k
        _sub_column_writers.push_back(std::move(sub_column_writer));
939
1.38k
    }
940
230
    _num_sub_column_writers = _sub_column_writers.size();
941
230
    DCHECK(_num_sub_column_writers >= 1);
942
230
    if (is_nullable()) {
943
230
        _null_writer.reset(null_writer);
944
230
    }
945
230
}
946
947
230
Status StructColumnWriter::init() {
948
1.38k
    for (auto& column_writer : _sub_column_writers) {
949
1.38k
        RETURN_IF_ERROR(column_writer->init());
950
1.38k
    }
951
230
    if (is_nullable()) {
952
230
        RETURN_IF_ERROR(_null_writer->init());
953
230
    }
954
230
    return Status::OK();
955
230
}
956
957
230
Status StructColumnWriter::write_inverted_index() {
958
230
    if (_opts.need_inverted_index) {
959
0
        for (auto& column_writer : _sub_column_writers) {
960
0
            RETURN_IF_ERROR(column_writer->write_inverted_index());
961
0
        }
962
0
    }
963
230
    return Status::OK();
964
230
}
965
966
Status StructColumnWriter::append_nullable(const uint8_t* null_map, const uint8_t** ptr,
967
230
                                           size_t num_rows) {
968
230
    RETURN_IF_ERROR(append_data(ptr, num_rows));
969
230
    RETURN_IF_ERROR(_null_writer->append_data(&null_map, num_rows));
970
230
    return Status::OK();
971
230
}
972
973
230
Status StructColumnWriter::append_data(const uint8_t** ptr, size_t num_rows) {
974
230
    const auto* results = reinterpret_cast<const uint64_t*>(*ptr);
975
1.61k
    for (size_t i = 0; i < _num_sub_column_writers; ++i) {
976
1.38k
        auto nullmap = *(results + _num_sub_column_writers + i);
977
1.38k
        auto data = *(results + i);
978
1.38k
        RETURN_IF_ERROR(_sub_column_writers[i]->append(reinterpret_cast<const uint8_t*>(nullmap),
979
1.38k
                                                       reinterpret_cast<const void*>(data),
980
1.38k
                                                       num_rows));
981
1.38k
    }
982
230
    return Status::OK();
983
230
}
984
985
0
uint64_t StructColumnWriter::estimate_buffer_size() {
986
0
    uint64_t size = 0;
987
0
    for (auto& column_writer : _sub_column_writers) {
988
0
        size += column_writer->estimate_buffer_size();
989
0
    }
990
0
    size += is_nullable() ? _null_writer->estimate_buffer_size() : 0;
991
0
    return size;
992
0
}
993
994
230
Status StructColumnWriter::finish() {
995
1.38k
    for (auto& column_writer : _sub_column_writers) {
996
1.38k
        RETURN_IF_ERROR(column_writer->finish());
997
1.38k
    }
998
230
    if (is_nullable()) {
999
230
        RETURN_IF_ERROR(_null_writer->finish());
1000
230
    }
1001
230
    _opts.meta->set_num_rows(get_next_rowid());
1002
230
    return Status::OK();
1003
230
}
1004
1005
230
Status StructColumnWriter::write_data() {
1006
1.38k
    for (auto& column_writer : _sub_column_writers) {
1007
1.38k
        RETURN_IF_ERROR(column_writer->write_data());
1008
1.38k
    }
1009
230
    if (is_nullable()) {
1010
230
        RETURN_IF_ERROR(_null_writer->write_data());
1011
230
    }
1012
230
    return Status::OK();
1013
230
}
1014
1015
230
Status StructColumnWriter::write_ordinal_index() {
1016
1.38k
    for (auto& column_writer : _sub_column_writers) {
1017
1.38k
        RETURN_IF_ERROR(column_writer->write_ordinal_index());
1018
1.38k
    }
1019
230
    if (is_nullable()) {
1020
230
        RETURN_IF_ERROR(_null_writer->write_ordinal_index());
1021
230
    }
1022
230
    return Status::OK();
1023
230
}
1024
1025
0
Status StructColumnWriter::append_nulls(size_t num_rows) {
1026
0
    for (auto& column_writer : _sub_column_writers) {
1027
0
        RETURN_IF_ERROR(column_writer->append_nulls(num_rows));
1028
0
    }
1029
0
    if (is_nullable()) {
1030
0
        std::vector<UInt8> null_signs(num_rows, 1);
1031
0
        const uint8_t* null_sign_ptr = null_signs.data();
1032
0
        RETURN_IF_ERROR(_null_writer->append_data(&null_sign_ptr, num_rows));
1033
0
    }
1034
0
    return Status::OK();
1035
0
}
1036
1037
0
Status StructColumnWriter::finish_current_page() {
1038
0
    return Status::NotSupported("struct writer has no data, can not finish_current_page");
1039
0
}
1040
1041
ArrayColumnWriter::ArrayColumnWriter(const ColumnWriterOptions& opts,
1042
                                     std::unique_ptr<StorageField> field,
1043
                                     OffsetColumnWriter* offset_writer,
1044
                                     ScalarColumnWriter* null_writer,
1045
                                     std::unique_ptr<ColumnWriter> item_writer)
1046
2.06k
        : ColumnWriter(std::move(field), opts.meta->is_nullable(), opts.meta),
1047
2.06k
          _item_writer(std::move(item_writer)),
1048
2.06k
          _opts(opts) {
1049
2.06k
    _offset_writer.reset(offset_writer);
1050
2.06k
    if (is_nullable()) {
1051
1.87k
        _null_writer.reset(null_writer);
1052
1.87k
    }
1053
2.06k
}
1054
1055
2.06k
Status ArrayColumnWriter::init() {
1056
2.06k
    RETURN_IF_ERROR(_offset_writer->init());
1057
2.06k
    if (is_nullable()) {
1058
1.87k
        RETURN_IF_ERROR(_null_writer->init());
1059
1.87k
    }
1060
2.06k
    RETURN_IF_ERROR(_item_writer->init());
1061
2.06k
    if (_opts.need_inverted_index) {
1062
0
        auto* writer = dynamic_cast<ScalarColumnWriter*>(_item_writer.get());
1063
0
        if (writer != nullptr) {
1064
0
            RETURN_IF_ERROR(IndexColumnWriter::create(get_field(), &_inverted_index_writer,
1065
0
                                                      _opts.index_file_writer,
1066
0
                                                      _opts.inverted_indexes[0]));
1067
0
        }
1068
0
    }
1069
2.06k
    if (_opts.need_ann_index) {
1070
1
        auto* writer = dynamic_cast<ScalarColumnWriter*>(_item_writer.get());
1071
1
        if (writer != nullptr) {
1072
1
            _ann_index_writer = std::make_unique<AnnIndexColumnWriter>(_opts.index_file_writer,
1073
1
                                                                       _opts.ann_index);
1074
1
            RETURN_IF_ERROR(_ann_index_writer->init());
1075
1
        }
1076
1
    }
1077
2.06k
    return Status::OK();
1078
2.06k
}
1079
1080
1.94k
Status ArrayColumnWriter::write_inverted_index() {
1081
1.94k
    if (_opts.need_inverted_index) {
1082
0
        return _inverted_index_writer->finish();
1083
0
    }
1084
1.94k
    return Status::OK();
1085
1.94k
}
1086
1087
1.94k
Status ArrayColumnWriter::write_ann_index() {
1088
1.94k
    if (_opts.need_ann_index) {
1089
1
        return _ann_index_writer->finish();
1090
1
    }
1091
1.94k
    return Status::OK();
1092
1.94k
}
1093
1094
// batch append data for array
1095
2.12k
Status ArrayColumnWriter::append_data(const uint8_t** ptr, size_t num_rows) {
1096
    // data_ptr contains
1097
    // [size, offset_ptr, item_data_ptr, item_nullmap_ptr]
1098
2.12k
    auto data_ptr = reinterpret_cast<const uint64_t*>(*ptr);
1099
    // total number length
1100
2.12k
    size_t element_cnt = size_t((unsigned long)(*data_ptr));
1101
2.12k
    auto offset_data = *(data_ptr + 1);
1102
2.12k
    const uint8_t* offsets_ptr = (const uint8_t*)offset_data;
1103
2.12k
    auto data = *(data_ptr + 2);
1104
2.12k
    auto nested_null_map = *(data_ptr + 3);
1105
2.12k
    if (element_cnt > 0) {
1106
1.59k
        RETURN_IF_ERROR(_item_writer->append(reinterpret_cast<const uint8_t*>(nested_null_map),
1107
1.59k
                                             reinterpret_cast<const void*>(data), element_cnt));
1108
1.59k
    }
1109
2.12k
    if (_opts.need_inverted_index) {
1110
0
        auto* writer = dynamic_cast<ScalarColumnWriter*>(_item_writer.get());
1111
        // now only support nested type is scala
1112
0
        if (writer != nullptr) {
1113
            //NOTE: use array field name as index field, but item_writer size should be used when moving item_data_ptr
1114
0
            RETURN_IF_ERROR(_inverted_index_writer->add_array_values(
1115
0
                    _item_writer->get_field()->size(), reinterpret_cast<const void*>(data),
1116
0
                    reinterpret_cast<const uint8_t*>(nested_null_map), offsets_ptr, num_rows));
1117
0
        }
1118
0
    }
1119
1120
2.12k
    if (_opts.need_ann_index) {
1121
1
        auto* writer = dynamic_cast<ScalarColumnWriter*>(_item_writer.get());
1122
        // now only support nested type is scala
1123
1
        if (writer != nullptr) {
1124
            //NOTE: use array field name as index field, but item_writer size should be used when moving item_data_ptr
1125
1
            RETURN_IF_ERROR(_ann_index_writer->add_array_values(
1126
1
                    _item_writer->get_field()->size(), reinterpret_cast<const void*>(data),
1127
1
                    reinterpret_cast<const uint8_t*>(nested_null_map), offsets_ptr, num_rows));
1128
1
        } else {
1129
0
            return Status::NotSupported(
1130
0
                    "Ann index can only be build on array with scalar type. but got {} as "
1131
0
                    "nested",
1132
0
                    _item_writer->get_field()->type());
1133
0
        }
1134
1
    }
1135
1136
2.12k
    RETURN_IF_ERROR(_offset_writer->append_data(&offsets_ptr, num_rows));
1137
2.12k
    return Status::OK();
1138
2.12k
}
1139
1140
659
uint64_t ArrayColumnWriter::estimate_buffer_size() {
1141
659
    return _offset_writer->estimate_buffer_size() +
1142
659
           (is_nullable() ? _null_writer->estimate_buffer_size() : 0) +
1143
659
           _item_writer->estimate_buffer_size();
1144
659
}
1145
1146
Status ArrayColumnWriter::append_nullable(const uint8_t* null_map, const uint8_t** ptr,
1147
1.93k
                                          size_t num_rows) {
1148
1.93k
    RETURN_IF_ERROR(append_data(ptr, num_rows));
1149
1.93k
    if (is_nullable()) {
1150
1.93k
        if (_opts.need_inverted_index) {
1151
0
            RETURN_IF_ERROR(_inverted_index_writer->add_array_nulls(null_map, num_rows));
1152
0
        }
1153
1.93k
        RETURN_IF_ERROR(_null_writer->append_data(&null_map, num_rows));
1154
1.93k
    }
1155
1.93k
    return Status::OK();
1156
1.93k
}
1157
1158
2.06k
Status ArrayColumnWriter::finish() {
1159
2.06k
    RETURN_IF_ERROR(_offset_writer->finish());
1160
2.06k
    if (is_nullable()) {
1161
1.87k
        RETURN_IF_ERROR(_null_writer->finish());
1162
1.87k
    }
1163
2.06k
    RETURN_IF_ERROR(_item_writer->finish());
1164
2.06k
    _opts.meta->set_num_rows(get_next_rowid());
1165
2.06k
    return Status::OK();
1166
2.06k
}
1167
1168
2.06k
Status ArrayColumnWriter::write_data() {
1169
2.06k
    RETURN_IF_ERROR(_offset_writer->write_data());
1170
2.06k
    if (is_nullable()) {
1171
1.87k
        RETURN_IF_ERROR(_null_writer->write_data());
1172
1.87k
    }
1173
2.06k
    RETURN_IF_ERROR(_item_writer->write_data());
1174
2.06k
    return Status::OK();
1175
2.06k
}
1176
1177
2.03k
Status ArrayColumnWriter::write_ordinal_index() {
1178
2.03k
    RETURN_IF_ERROR(_offset_writer->write_ordinal_index());
1179
2.03k
    if (is_nullable()) {
1180
1.84k
        RETURN_IF_ERROR(_null_writer->write_ordinal_index());
1181
1.84k
    }
1182
2.03k
    if (!has_empty_items()) {
1183
1.53k
        RETURN_IF_ERROR(_item_writer->write_ordinal_index());
1184
1.53k
    }
1185
2.03k
    return Status::OK();
1186
2.03k
}
1187
1188
0
Status ArrayColumnWriter::append_nulls(size_t num_rows) {
1189
0
    const UInt64 offset = cast_set<UInt64>(_item_writer->get_next_rowid());
1190
0
    std::vector<UInt64> offsets_data(num_rows + 1, offset);
1191
0
    const uint8_t* offsets_ptr = reinterpret_cast<const uint8_t*>(offsets_data.data());
1192
0
    RETURN_IF_ERROR(_offset_writer->append_data(&offsets_ptr, num_rows));
1193
0
    return write_null_column(num_rows, true);
1194
0
}
1195
1196
0
Status ArrayColumnWriter::write_null_column(size_t num_rows, bool is_null) {
1197
0
    uint8_t null_sign = is_null ? 1 : 0;
1198
0
    while (is_nullable() && num_rows > 0) {
1199
        // TODO llj bulk write
1200
0
        const uint8_t* null_sign_ptr = &null_sign;
1201
0
        RETURN_IF_ERROR(_null_writer->append_data(&null_sign_ptr, 1));
1202
0
        --num_rows;
1203
0
    }
1204
0
    return Status::OK();
1205
0
}
1206
1207
0
Status ArrayColumnWriter::finish_current_page() {
1208
0
    return Status::NotSupported("array writer has no data, can not finish_current_page");
1209
0
}
1210
1211
/// ============================= MapColumnWriter =====================////
1212
MapColumnWriter::MapColumnWriter(const ColumnWriterOptions& opts,
1213
                                 std::unique_ptr<StorageField> field,
1214
                                 ScalarColumnWriter* null_writer, OffsetColumnWriter* offset_writer,
1215
                                 std::vector<std::unique_ptr<ColumnWriter>>& kv_writers)
1216
1.67k
        : ColumnWriter(std::move(field), opts.meta->is_nullable(), opts.meta), _opts(opts) {
1217
1.67k
    CHECK_EQ(kv_writers.size(), 2);
1218
1.67k
    _offsets_writer.reset(offset_writer);
1219
1.67k
    if (is_nullable()) {
1220
1.35k
        _null_writer.reset(null_writer);
1221
1.35k
    }
1222
3.35k
    for (auto& sub_writers : kv_writers) {
1223
3.35k
        _kv_writers.push_back(std::move(sub_writers));
1224
3.35k
    }
1225
1.67k
}
1226
1227
1.67k
Status MapColumnWriter::init() {
1228
1.67k
    RETURN_IF_ERROR(_offsets_writer->init());
1229
1.67k
    if (is_nullable()) {
1230
1.35k
        RETURN_IF_ERROR(_null_writer->init());
1231
1.35k
    }
1232
    // here register_flush_page_callback to call this.put_extra_info_in_page()
1233
    // when finish cur data page
1234
3.35k
    for (auto& sub_writer : _kv_writers) {
1235
3.35k
        RETURN_IF_ERROR(sub_writer->init());
1236
3.35k
    }
1237
1.67k
    return Status::OK();
1238
1.67k
}
1239
1240
977
uint64_t MapColumnWriter::estimate_buffer_size() {
1241
977
    size_t estimate = 0;
1242
1.95k
    for (auto& sub_writer : _kv_writers) {
1243
1.95k
        estimate += sub_writer->estimate_buffer_size();
1244
1.95k
    }
1245
977
    estimate += _offsets_writer->estimate_buffer_size();
1246
977
    if (is_nullable()) {
1247
976
        estimate += _null_writer->estimate_buffer_size();
1248
976
    }
1249
977
    return estimate;
1250
977
}
1251
1252
1.67k
Status MapColumnWriter::finish() {
1253
1.67k
    RETURN_IF_ERROR(_offsets_writer->finish());
1254
1.67k
    if (is_nullable()) {
1255
1.35k
        RETURN_IF_ERROR(_null_writer->finish());
1256
1.35k
    }
1257
3.34k
    for (auto& sub_writer : _kv_writers) {
1258
3.34k
        RETURN_IF_ERROR(sub_writer->finish());
1259
3.34k
    }
1260
1.67k
    _opts.meta->set_num_rows(get_next_rowid());
1261
1.67k
    return Status::OK();
1262
1.67k
}
1263
1264
Status MapColumnWriter::append_nullable(const uint8_t* null_map, const uint8_t** ptr,
1265
1.52k
                                        size_t num_rows) {
1266
1.52k
    RETURN_IF_ERROR(append_data(ptr, num_rows));
1267
1.52k
    if (is_nullable()) {
1268
1.52k
        RETURN_IF_ERROR(_null_writer->append_data(&null_map, num_rows));
1269
1.52k
    }
1270
1.52k
    return Status::OK();
1271
1.52k
}
1272
1273
// write key value data with offsets
1274
1.85k
Status MapColumnWriter::append_data(const uint8_t** ptr, size_t num_rows) {
1275
    // data_ptr contains
1276
    // [size, offset_ptr, key_data_ptr, val_data_ptr, k_nullmap_ptr, v_nullmap_pr]
1277
    // which converted results from olap_map_convertor and later will use a structure to replace it
1278
1.85k
    auto data_ptr = reinterpret_cast<const uint64_t*>(*ptr);
1279
    // total number length
1280
1.85k
    size_t element_cnt = size_t((unsigned long)(*data_ptr));
1281
1.85k
    auto offset_data = *(data_ptr + 1);
1282
1.85k
    const uint8_t* offsets_ptr = (const uint8_t*)offset_data;
1283
1284
1.85k
    if (element_cnt > 0) {
1285
4.90k
        for (size_t i = 0; i < 2; ++i) {
1286
3.27k
            auto data = *(data_ptr + 2 + i);
1287
3.27k
            auto nested_null_map = *(data_ptr + 2 + 2 + i);
1288
3.27k
            RETURN_IF_ERROR(
1289
3.27k
                    _kv_writers[i]->append(reinterpret_cast<const uint8_t*>(nested_null_map),
1290
3.27k
                                           reinterpret_cast<const void*>(data), element_cnt));
1291
3.27k
        }
1292
1.63k
    }
1293
    // make sure the order : offset writer flush next_array_item_ordinal after kv_writers append_data
1294
    // because we use _kv_writers[0]->get_next_rowid() to set next_array_item_ordinal in offset page footer
1295
1.85k
    RETURN_IF_ERROR(_offsets_writer->append_data(&offsets_ptr, num_rows));
1296
1.85k
    return Status::OK();
1297
1.85k
}
1298
1299
1.67k
Status MapColumnWriter::write_data() {
1300
1.67k
    RETURN_IF_ERROR(_offsets_writer->write_data());
1301
1.67k
    if (is_nullable()) {
1302
1.35k
        RETURN_IF_ERROR(_null_writer->write_data());
1303
1.35k
    }
1304
3.34k
    for (auto& sub_writer : _kv_writers) {
1305
3.34k
        RETURN_IF_ERROR(sub_writer->write_data());
1306
3.34k
    }
1307
1.67k
    return Status::OK();
1308
1.67k
}
1309
1310
1.67k
Status MapColumnWriter::write_ordinal_index() {
1311
1.67k
    RETURN_IF_ERROR(_offsets_writer->write_ordinal_index());
1312
1.67k
    if (is_nullable()) {
1313
1.35k
        RETURN_IF_ERROR(_null_writer->write_ordinal_index());
1314
1.35k
    }
1315
3.34k
    for (auto& sub_writer : _kv_writers) {
1316
3.34k
        if (sub_writer->get_next_rowid() != 0) {
1317
2.91k
            RETURN_IF_ERROR(sub_writer->write_ordinal_index());
1318
2.91k
        }
1319
3.34k
    }
1320
1.67k
    return Status::OK();
1321
1.67k
}
1322
1323
0
Status MapColumnWriter::append_nulls(size_t num_rows) {
1324
0
    for (auto& sub_writer : _kv_writers) {
1325
0
        RETURN_IF_ERROR(sub_writer->append_nulls(num_rows));
1326
0
    }
1327
0
    const UInt64 offset = cast_set<UInt64>(_kv_writers[0]->get_next_rowid());
1328
0
    std::vector<UInt64> offsets_data(num_rows + 1, offset);
1329
0
    const uint8_t* offsets_ptr = reinterpret_cast<const uint8_t*>(offsets_data.data());
1330
0
    RETURN_IF_ERROR(_offsets_writer->append_data(&offsets_ptr, num_rows));
1331
1332
0
    if (is_nullable()) {
1333
0
        std::vector<UInt8> null_signs(num_rows, 1);
1334
0
        const uint8_t* null_sign_ptr = null_signs.data();
1335
0
        RETURN_IF_ERROR(_null_writer->append_data(&null_sign_ptr, num_rows));
1336
0
    }
1337
0
    return Status::OK();
1338
0
}
1339
1340
0
Status MapColumnWriter::finish_current_page() {
1341
0
    return Status::NotSupported("map writer has no data, can not finish_current_page");
1342
0
}
1343
1344
1.34k
Status MapColumnWriter::write_inverted_index() {
1345
1.34k
    if (_opts.need_inverted_index) {
1346
0
        return _index_builder->finish();
1347
0
    }
1348
1.34k
    return Status::OK();
1349
1.34k
}
1350
1351
VariantColumnWriter::VariantColumnWriter(const ColumnWriterOptions& opts,
1352
                                         const TabletColumn* column,
1353
                                         std::unique_ptr<StorageField> field)
1354
316
        : ColumnWriter(std::move(field), opts.meta->is_nullable(), opts.meta) {
1355
316
    _impl = std::make_unique<VariantColumnWriterImpl>(opts, column);
1356
316
}
1357
1358
316
Status VariantColumnWriter::init() {
1359
316
    return _impl->init();
1360
316
}
1361
1362
324
Status VariantColumnWriter::append_data(const uint8_t** ptr, size_t num_rows) {
1363
324
    _next_rowid += num_rows;
1364
324
    return _impl->append_data(ptr, num_rows);
1365
324
}
1366
1367
769
uint64_t VariantColumnWriter::estimate_buffer_size() {
1368
769
    return _impl->estimate_buffer_size();
1369
769
}
1370
1371
311
Status VariantColumnWriter::finish() {
1372
311
    return _impl->finish();
1373
311
}
1374
312
Status VariantColumnWriter::write_data() {
1375
312
    return _impl->write_data();
1376
312
}
1377
311
Status VariantColumnWriter::write_ordinal_index() {
1378
311
    return _impl->write_ordinal_index();
1379
311
}
1380
1381
308
Status VariantColumnWriter::write_zone_map() {
1382
308
    return _impl->write_zone_map();
1383
308
}
1384
1385
292
Status VariantColumnWriter::write_inverted_index() {
1386
292
    return _impl->write_inverted_index();
1387
292
}
1388
293
Status VariantColumnWriter::write_bloom_filter_index() {
1389
293
    return _impl->write_bloom_filter_index();
1390
293
}
1391
1392
Status VariantColumnWriter::append_nullable(const uint8_t* null_map, const uint8_t** ptr,
1393
13
                                            size_t num_rows) {
1394
13
    return _impl->append_nullable(null_map, ptr, num_rows);
1395
13
}
1396
1397
} // namespace doris::segment_v2