Coverage Report

Created: 2026-03-13 19:41

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