Coverage Report

Created: 2026-05-29 17:33

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
be/src/storage/segment/segment_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/segment_writer.h"
19
20
#include <assert.h>
21
#include <gen_cpp/segment_v2.pb.h>
22
#include <parallel_hashmap/phmap.h>
23
24
#include <algorithm>
25
26
// IWYU pragma: no_include <opentelemetry/common/threadlocal.h>
27
#include <crc32c/crc32c.h>
28
29
#include "cloud/config.h"
30
#include "common/cast_set.h"
31
#include "common/compiler_util.h" // IWYU pragma: keep
32
#include "common/config.h"
33
#include "common/logging.h" // LOG
34
#include "common/status.h"
35
#include "core/block/block.h"
36
#include "core/block/column_with_type_and_name.h"
37
#include "core/column/column_nullable.h"
38
#include "core/data_type/primitive_type.h"
39
#include "core/field.h"
40
#include "core/types.h"
41
#include "core/value/vdatetime_value.h"
42
#include "exec/common/variant_util.h"
43
#include "io/cache/block_file_cache.h"
44
#include "io/cache/block_file_cache_factory.h"
45
#include "io/fs/file_system.h"
46
#include "io/fs/file_writer.h"
47
#include "io/fs/local_file_system.h"
48
#include "runtime/exec_env.h"
49
#include "runtime/memory/mem_tracker.h"
50
#include "service/point_query_executor.h"
51
#include "storage/data_dir.h"
52
#include "storage/index/index_file_writer.h"
53
#include "storage/index/index_writer.h"
54
#include "storage/index/inverted/inverted_index_fs_directory.h"
55
#include "storage/index/primary_key_index.h"
56
#include "storage/index/short_key_index.h"
57
#include "storage/iterator/olap_data_convertor.h"
58
#include "storage/key_coder.h"
59
#include "storage/olap_common.h"
60
#include "storage/olap_define.h"
61
#include "storage/partial_update_info.h"
62
#include "storage/rowset/rowset_writer_context.h" // RowsetWriterContext
63
#include "storage/rowset/segment_creator.h"
64
#include "storage/segment/column_writer.h" // ColumnWriter
65
#include "storage/segment/encoding_info.h"
66
#include "storage/segment/external_col_meta_util.h"
67
#include "storage/segment/historical_row_retriever.h"
68
#include "storage/segment/page_io.h"
69
#include "storage/segment/page_pointer.h"
70
#include "storage/segment/segment_loader.h"
71
#include "storage/segment/variant/variant_ext_meta_writer.h"
72
#include "storage/segment/variant_stats_calculator.h"
73
#include "storage/storage_engine.h"
74
#include "storage/tablet/tablet_schema.h"
75
#include "storage/utils.h"
76
#include "util/coding.h"
77
#include "util/faststring.h"
78
#include "util/jsonb/serialize.h"
79
#include "util/simd/bits.h"
80
namespace doris {
81
namespace segment_v2 {
82
83
using namespace ErrorCode;
84
using namespace KeyConsts;
85
86
const char* k_segment_magic = "D0R1";
87
const uint32_t k_segment_magic_length = 4;
88
89
2.44k
inline std::string segment_mem_tracker_name(uint32_t segment_id) {
90
2.44k
    return "SegmentWriter:Segment-" + std::to_string(segment_id);
91
2.44k
}
92
93
SegmentWriter::SegmentWriter(io::FileWriter* file_writer, uint32_t segment_id,
94
                             TabletSchemaSPtr tablet_schema, BaseTabletSPtr tablet,
95
                             DataDir* data_dir, const SegmentWriterOptions& opts,
96
                             IndexFileWriter* index_file_writer)
97
2.44k
        : _segment_id(segment_id),
98
2.44k
          _tablet_schema(std::move(tablet_schema)),
99
2.44k
          _tablet(std::move(tablet)),
100
2.44k
          _data_dir(data_dir),
101
2.44k
          _opts(opts),
102
2.44k
          _file_writer(file_writer),
103
2.44k
          _index_file_writer(index_file_writer),
104
2.44k
          _mem_tracker(std::make_unique<MemTracker>(segment_mem_tracker_name(segment_id))),
105
2.44k
          _mow_context(std::move(opts.mow_ctx)) {
106
2.44k
    CHECK_NOTNULL(file_writer);
107
2.44k
    _num_sort_key_columns = _tablet_schema->num_key_columns();
108
2.44k
    _num_short_key_columns = _tablet_schema->num_short_key_columns();
109
2.44k
    if (!_is_mow_with_cluster_key()) {
110
2.44k
        DCHECK(_num_sort_key_columns >= _num_short_key_columns)
111
0
                << ", table_id=" << _tablet_schema->table_id()
112
0
                << ", num_key_columns=" << _num_sort_key_columns
113
0
                << ", num_short_key_columns=" << _num_short_key_columns
114
0
                << ", cluster_key_columns=" << _tablet_schema->cluster_key_uids().size();
115
2.44k
    }
116
4.50k
    for (size_t cid = 0; cid < _num_sort_key_columns; ++cid) {
117
2.05k
        const auto& column = _tablet_schema->column(cid);
118
2.05k
        _key_coders.push_back(get_key_coder(column.type()));
119
2.05k
        _key_index_size.push_back(cast_set<uint16_t>(column.index_length()));
120
2.05k
    }
121
2.44k
    if (_is_mow()) {
122
        // encode the sequence id into the primary key index
123
69
        if (_tablet_schema->has_sequence_col()) {
124
26
            const auto& column = _tablet_schema->column(_tablet_schema->sequence_col_idx());
125
26
            _seq_coder = get_key_coder(column.type());
126
26
        }
127
        // encode the rowid into the primary key index
128
69
        if (_is_mow_with_cluster_key()) {
129
0
            _rowid_coder = get_key_coder(FieldType::OLAP_FIELD_TYPE_UNSIGNED_INT);
130
            // primary keys
131
0
            _primary_key_coders.swap(_key_coders);
132
            // cluster keys
133
0
            _key_coders.clear();
134
0
            _key_index_size.clear();
135
0
            _num_sort_key_columns = _tablet_schema->cluster_key_uids().size();
136
0
            for (auto cid : _tablet_schema->cluster_key_uids()) {
137
0
                const auto& column = _tablet_schema->column_by_uid(cid);
138
0
                _key_coders.push_back(get_key_coder(column.type()));
139
0
                _key_index_size.push_back(cast_set<uint16_t>(column.index_length()));
140
0
            }
141
0
        }
142
69
    }
143
2.44k
}
144
145
2.44k
SegmentWriter::~SegmentWriter() {
146
2.44k
    _mem_tracker->release(_mem_tracker->consumption());
147
2.44k
}
148
149
void SegmentWriter::init_column_meta(ColumnMetaPB* meta, uint32_t column_id,
150
10.2k
                                     const TabletColumn& column, const ColumnWriterOptions& opts) {
151
10.2k
    meta->set_column_id(column_id);
152
10.2k
    meta->set_type(int(column.type()));
153
10.2k
    meta->set_length(column.length());
154
10.2k
    meta->set_encoding(EncodingInfo::resolve_default_encoding(opts.storage_format, column));
155
10.2k
    meta->set_compression(_opts.compression_type);
156
10.2k
    meta->set_is_nullable(column.is_nullable());
157
10.2k
    meta->set_default_value(column.default_value());
158
10.2k
    meta->set_precision(column.precision());
159
10.2k
    meta->set_frac(column.frac());
160
10.2k
    if (column.has_path_info()) {
161
288
        column.path_info_ptr()->to_protobuf(meta->mutable_column_path_info(),
162
288
                                            column.parent_unique_id());
163
288
    }
164
10.2k
    meta->set_unique_id(column.unique_id());
165
10.3k
    for (uint32_t i = 0; i < column.get_subtype_count(); ++i) {
166
15
        init_column_meta(meta->add_children_columns(), column_id, column.get_sub_column(i), opts);
167
15
    }
168
10.2k
    meta->set_result_is_nullable(column.get_result_is_nullable());
169
10.2k
    meta->set_function_name(column.get_aggregation_name());
170
10.2k
    meta->set_be_exec_version(column.get_be_exec_version());
171
10.2k
    if (column.is_variant_type()) {
172
285
        meta->set_variant_max_subcolumns_count(column.variant_max_subcolumns_count());
173
285
        meta->set_variant_enable_doc_mode(column.variant_enable_doc_mode());
174
285
    }
175
10.2k
}
176
177
1.25k
Status SegmentWriter::init() {
178
1.25k
    std::vector<uint32_t> column_ids;
179
1.25k
    auto column_cnt = cast_set<int>(_tablet_schema->num_columns());
180
8.13k
    for (uint32_t i = 0; i < column_cnt; ++i) {
181
6.87k
        column_ids.emplace_back(i);
182
6.87k
    }
183
1.25k
    return init(column_ids, true);
184
1.25k
}
185
186
Status SegmentWriter::_create_column_writer(uint32_t cid, const TabletColumn& column,
187
10.2k
                                            const TabletSchemaSPtr& schema) {
188
10.2k
    ColumnWriterOptions opts;
189
10.2k
    opts.meta = _footer.add_columns();
190
10.2k
    opts.storage_format = schema->storage_format();
191
192
10.2k
    init_column_meta(opts.meta, cid, column, opts);
193
194
    // now we create zone map for key columns in AGG_KEYS or all column in UNIQUE_KEYS or DUP_KEYS
195
    // except for columns whose type don't support zone map.
196
10.2k
    opts.need_zone_map = column.is_key() || schema->keys_type() != KeysType::AGG_KEYS;
197
10.2k
    opts.need_bloom_filter = column.is_bf_column();
198
10.2k
    if (opts.need_bloom_filter) {
199
4
        opts.bf_options.fpp = schema->has_bf_fpp() ? schema->bloom_filter_fpp() : 0.05;
200
4
    }
201
10.2k
    auto* tablet_index = schema->get_ngram_bf_index(column.unique_id());
202
10.2k
    if (tablet_index) {
203
0
        opts.need_bloom_filter = true;
204
0
        opts.is_ngram_bf_index = true;
205
        //narrow convert from int32_t to uint8_t and uint16_t which is dangerous
206
0
        auto gram_size = tablet_index->get_gram_size();
207
0
        auto gram_bf_size = tablet_index->get_gram_bf_size();
208
0
        if (gram_size > 256 || gram_size < 1) {
209
0
            return Status::NotSupported("Do not support ngram bloom filter for ngram_size: ",
210
0
                                        gram_size);
211
0
        }
212
0
        if (gram_bf_size > 65535 || gram_bf_size < 64) {
213
0
            return Status::NotSupported("Do not support ngram bloom filter for bf_size: ",
214
0
                                        gram_bf_size);
215
0
        }
216
0
        opts.gram_size = cast_set<uint8_t>(gram_size);
217
0
        opts.gram_bf_size = cast_set<uint16_t>(gram_bf_size);
218
0
    }
219
220
10.2k
    bool skip_inverted_index = false;
221
10.2k
    if (_opts.rowset_ctx != nullptr) {
222
        // skip write inverted index for index compaction column
223
8.07k
        skip_inverted_index =
224
8.07k
                _opts.rowset_ctx->columns_to_do_index_compaction.count(column.unique_id()) > 0;
225
8.07k
    }
226
    // skip write inverted index on load if skip_write_index_on_load is true
227
10.2k
    if (_opts.write_type == DataWriteType::TYPE_DIRECT && schema->skip_write_index_on_load()) {
228
0
        skip_inverted_index = true;
229
0
    }
230
    // indexes for this column
231
10.2k
    if (!skip_inverted_index) {
232
9.85k
        auto inverted_indexs = schema->inverted_indexs(column);
233
9.85k
        if (!inverted_indexs.empty()) {
234
2.18k
            opts.inverted_indexes = inverted_indexs;
235
2.18k
            opts.need_inverted_index = true;
236
2.18k
            DCHECK(_index_file_writer != nullptr);
237
2.18k
        }
238
9.85k
    }
239
    // indexes for this column
240
10.2k
    if (const auto& index = schema->ann_index(column); index != nullptr) {
241
1
        opts.ann_index = index;
242
1
        opts.need_ann_index = true;
243
1
        DCHECK(_index_file_writer != nullptr);
244
1
    }
245
246
10.2k
    opts.index_file_writer = _index_file_writer;
247
248
10.2k
#define DISABLE_INDEX_IF_FIELD_TYPE(TYPE)                     \
249
92.5k
    if (column.type() == FieldType::OLAP_FIELD_TYPE_##TYPE) { \
250
289
        opts.need_zone_map = false;                           \
251
289
        opts.need_bloom_filter = false;                       \
252
289
    }
253
254
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(STRUCT)
255
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(ARRAY)
256
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(JSONB)
257
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(AGG_STATE)
258
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(MAP)
259
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(BITMAP)
260
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(HLL)
261
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(QUANTILE_STATE)
262
10.2k
    DISABLE_INDEX_IF_FIELD_TYPE(VARIANT)
263
264
10.2k
#undef DISABLE_INDEX_IF_FIELD_TYPE
265
266
10.2k
    int64_t storage_page_size = _tablet_schema->storage_page_size();
267
    // storage_page_size must be between 4KB and 10MB.
268
10.2k
    if (storage_page_size >= 4096 && storage_page_size <= 10485760) {
269
10.2k
        opts.data_page_size = storage_page_size;
270
10.2k
    }
271
10.2k
    opts.dict_page_size = _tablet_schema->storage_dict_page_size();
272
10.2k
    DBUG_EXECUTE_IF("VerticalSegmentWriter._create_column_writer.storage_page_size", {
273
10.2k
        auto table_id = DebugPoints::instance()->get_debug_param_or_default<int64_t>(
274
10.2k
                "VerticalSegmentWriter._create_column_writer.storage_page_size", "table_id",
275
10.2k
                INT_MIN);
276
10.2k
        auto target_data_page_size = DebugPoints::instance()->get_debug_param_or_default<int64_t>(
277
10.2k
                "VerticalSegmentWriter._create_column_writer.storage_page_size",
278
10.2k
                "storage_page_size", INT_MIN);
279
10.2k
        if (table_id == INT_MIN || target_data_page_size == INT_MIN) {
280
10.2k
            return Status::Error<ErrorCode::INTERNAL_ERROR>(
281
10.2k
                    "Debug point parameters missing: either 'table_id' or 'storage_page_size' not "
282
10.2k
                    "set.");
283
10.2k
        }
284
10.2k
        if (table_id == _tablet_schema->table_id() &&
285
10.2k
            opts.data_page_size != target_data_page_size) {
286
10.2k
            return Status::Error<ErrorCode::INTERNAL_ERROR>(
287
10.2k
                    "Mismatch in 'storage_page_size': expected size does not match the current "
288
10.2k
                    "data page size. "
289
10.2k
                    "Expected: " +
290
10.2k
                    std::to_string(target_data_page_size) +
291
10.2k
                    ", Actual: " + std::to_string(opts.data_page_size) + ".");
292
10.2k
        }
293
10.2k
    })
294
10.2k
    if (column.is_row_store_column()) {
295
        // smaller page size for row store column; encoding is already set to PLAIN /
296
        // PLAIN_V2 by init_column_meta via resolve_default_encoding().
297
1
        auto page_size = _tablet_schema->row_store_page_size();
298
1
        opts.data_page_size =
299
1
                (page_size > 0) ? page_size : segment_v2::ROW_STORE_PAGE_SIZE_DEFAULT_VALUE;
300
1
    }
301
302
10.2k
    opts.rowset_ctx = _opts.rowset_ctx;
303
10.2k
    opts.file_writer = _file_writer;
304
10.2k
    opts.compression_type = _opts.compression_type;
305
10.2k
    opts.footer = &_footer;
306
10.2k
    if (_opts.rowset_ctx != nullptr) {
307
8.07k
        opts.input_rs_readers = _opts.rowset_ctx->input_rs_readers;
308
8.07k
    }
309
310
10.2k
    std::unique_ptr<ColumnWriter> writer;
311
10.2k
    RETURN_IF_ERROR(ColumnWriter::create(opts, &column, _file_writer, &writer));
312
10.2k
    RETURN_IF_ERROR(writer->init());
313
10.2k
    _column_writers.push_back(std::move(writer));
314
315
10.2k
    _olap_data_convertor->add_column_data_convertor(column);
316
10.2k
    return Status::OK();
317
10.2k
}
318
319
3.37k
Status SegmentWriter::init(const std::vector<uint32_t>& col_ids, bool has_key) {
320
3.37k
    DCHECK(_column_writers.empty());
321
3.37k
    DCHECK(_column_ids.empty());
322
3.37k
    _has_key = has_key;
323
3.37k
    _column_writers.reserve(_tablet_schema->columns().size());
324
3.37k
    _column_ids.insert(_column_ids.end(), col_ids.begin(), col_ids.end());
325
3.37k
    _olap_data_convertor = std::make_unique<OlapBlockDataConvertor>();
326
3.37k
    if (_opts.compression_type == UNKNOWN_COMPRESSION) {
327
2.44k
        _opts.compression_type = _tablet_schema->compression_type();
328
2.44k
    }
329
330
    // Vertical compaction calls init() multiple times against the same writer; the footer accumulates entries
331
    // across calls, so this init()'s slice of footer columns starts at the current size.
332
3.37k
    const int variant_stats_footer_offset = _footer.columns_size();
333
3.37k
    RETURN_IF_ERROR(_create_writers(_tablet_schema, col_ids));
334
335
    // Initialize variant statistics calculator
336
3.37k
    _variant_stats_calculator = std::make_unique<VariantStatsCaculator>(
337
3.37k
            &_footer, _tablet_schema, col_ids, variant_stats_footer_offset);
338
339
    // we don't need the short key index for unique key merge on write table.
340
3.37k
    if (_has_key) {
341
2.44k
        if (_is_mow()) {
342
69
            size_t seq_col_length = 0;
343
69
            if (_tablet_schema->has_sequence_col()) {
344
26
                seq_col_length =
345
26
                        _tablet_schema->column(_tablet_schema->sequence_col_idx()).length() + 1;
346
26
            }
347
69
            size_t rowid_length = 0;
348
69
            if (_is_mow_with_cluster_key()) {
349
0
                rowid_length = PrimaryKeyIndexReader::ROW_ID_LENGTH;
350
0
                _short_key_index_builder.reset(
351
0
                        new ShortKeyIndexBuilder(_segment_id, _opts.num_rows_per_block));
352
0
            }
353
69
            _primary_key_index_builder.reset(
354
69
                    new PrimaryKeyIndexBuilder(_file_writer, seq_col_length, rowid_length));
355
69
            RETURN_IF_ERROR(_primary_key_index_builder->init());
356
2.37k
        } else {
357
2.37k
            _short_key_index_builder.reset(
358
2.37k
                    new ShortKeyIndexBuilder(_segment_id, _opts.num_rows_per_block));
359
2.37k
        }
360
2.44k
    }
361
3.37k
    return Status::OK();
362
3.37k
}
363
364
Status SegmentWriter::_create_writers(const TabletSchemaSPtr& tablet_schema,
365
3.37k
                                      const std::vector<uint32_t>& col_ids) {
366
3.37k
    _olap_data_convertor->reserve(col_ids.size());
367
10.2k
    for (auto& cid : col_ids) {
368
10.2k
        RETURN_IF_ERROR(_create_column_writer(cid, tablet_schema->column(cid), tablet_schema));
369
10.2k
    }
370
3.37k
    return Status::OK();
371
3.37k
}
372
373
800
void SegmentWriter::_maybe_invalid_row_cache(const std::string& key) {
374
    // Just invalid row cache for simplicity, since the rowset is not visible at present.
375
    // If we update/insert cache, if load failed rowset will not be visible but cached data
376
    // will be visible, and lead to inconsistency.
377
800
    if (!config::disable_storage_row_cache && _tablet_schema->has_row_store_for_all_columns() &&
378
800
        _opts.write_type == DataWriteType::TYPE_DIRECT) {
379
        // invalidate cache
380
0
        RowCache::instance()->erase({_opts.rowset_ctx->tablet_id, key});
381
0
    }
382
800
}
383
384
6
void SegmentWriter::_serialize_block_to_row_column(Block& block) {
385
6
    if (block.rows() == 0) {
386
0
        return;
387
0
    }
388
6
    MonotonicStopWatch watch;
389
6
    watch.start();
390
6
    int row_column_id = 0;
391
12
    for (int i = 0; i < _tablet_schema->num_columns(); ++i) {
392
6
        if (_tablet_schema->column(i).is_row_store_column()) {
393
0
            auto row_store_column_ptr = block.get_by_position(i).column->clone_empty();
394
0
            auto* row_store_column = static_cast<ColumnString*>(row_store_column_ptr.get());
395
0
            DataTypeSerDeSPtrs serdes = create_data_type_serdes(block.get_data_types());
396
0
            JsonbSerializeUtil::block_to_jsonb(*_tablet_schema, block, *row_store_column,
397
0
                                               cast_set<int>(_tablet_schema->num_columns()), serdes,
398
0
                                               {_tablet_schema->row_columns_uids().begin(),
399
0
                                                _tablet_schema->row_columns_uids().end()});
400
0
            block.replace_by_position(i, std::move(row_store_column_ptr));
401
0
            break;
402
0
        }
403
6
    }
404
405
6
    VLOG_DEBUG << "serialize , num_rows:" << block.rows() << ", row_column_id:" << row_column_id
406
0
               << ", total_byte_size:" << block.allocated_bytes() << ", serialize_cost(us)"
407
0
               << watch.elapsed_time() / 1000;
408
6
}
409
410
Status SegmentWriter::probe_key_for_mow(
411
        std::string key, std::size_t segment_pos, bool have_input_seq_column, bool have_delete_sign,
412
        const std::vector<RowsetSharedPtr>& specified_rowsets,
413
        std::vector<std::unique_ptr<SegmentCacheHandle>>& segment_caches,
414
        bool& has_default_or_nullable, std::vector<bool>& use_default_or_null_flag,
415
        const std::function<void(const RowLocation& loc)>& found_cb,
416
0
        const std::function<Status()>& not_found_cb, PartialUpdateStats& stats) {
417
0
    RowLocation loc;
418
    // save rowset shared ptr so this rowset wouldn't delete
419
0
    RowsetSharedPtr rowset;
420
0
    auto st = _tablet->lookup_row_key(
421
0
            key, _tablet_schema.get(), have_input_seq_column, specified_rowsets, &loc,
422
0
            cast_set<uint32_t>(_mow_context->max_version), segment_caches, &rowset);
423
0
    if (st.is<KEY_NOT_FOUND>()) {
424
0
        if (!have_delete_sign) {
425
0
            RETURN_IF_ERROR(not_found_cb());
426
0
        }
427
0
        ++stats.num_rows_new_added;
428
0
        has_default_or_nullable = true;
429
0
        use_default_or_null_flag.emplace_back(true);
430
0
        return Status::OK();
431
0
    }
432
0
    if (!st.ok() && !st.is<KEY_ALREADY_EXISTS>()) {
433
0
        LOG(WARNING) << "failed to lookup row key, error: " << st;
434
0
        return st;
435
0
    }
436
437
    // 1. if the delete sign is marked, it means that the value columns of the row will not
438
    //    be read. So we don't need to read the missing values from the previous rows.
439
    // 2. the one exception is when there are sequence columns in the table, we need to read
440
    //    the sequence columns, otherwise it may cause the merge-on-read based compaction
441
    //    policy to produce incorrect results
442
    // TODO(bobhan1): only read seq col rather than all columns in this situation for
443
    // partial update and flexible partial update
444
445
    // TODO(bobhan1): handle sequence column here
446
0
    if (st.is<KEY_ALREADY_EXISTS>() || (have_delete_sign && !_tablet_schema->has_sequence_col())) {
447
0
        has_default_or_nullable = true;
448
0
        use_default_or_null_flag.emplace_back(true);
449
0
    } else {
450
        // partial update should not contain invisible columns
451
0
        use_default_or_null_flag.emplace_back(false);
452
0
        _rsid_to_rowset.emplace(rowset->rowset_id(), rowset);
453
0
        found_cb(loc);
454
0
    }
455
456
0
    if (st.is<KEY_ALREADY_EXISTS>()) {
457
        // although we need to mark delete current row, we still need to read missing columns
458
        // for this row, we need to ensure that each column is aligned
459
0
        _mow_context->delete_bitmap->add(
460
0
                {_opts.rowset_ctx->rowset_id, _segment_id, DeleteBitmap::TEMP_VERSION_COMMON},
461
0
                cast_set<uint32_t>(segment_pos));
462
0
        ++stats.num_rows_deleted;
463
0
    } else {
464
0
        _mow_context->delete_bitmap->add(
465
0
                {loc.rowset_id, loc.segment_id, DeleteBitmap::TEMP_VERSION_COMMON}, loc.row_id);
466
0
        ++stats.num_rows_updated;
467
0
    }
468
0
    return Status::OK();
469
0
}
470
471
0
Status SegmentWriter::partial_update_preconditions_check(size_t row_pos) {
472
0
    if (!_is_mow()) {
473
0
        auto msg = fmt::format(
474
0
                "Can only do partial update on merge-on-write unique table, but found: "
475
0
                "keys_type={}, _opts.enable_unique_key_merge_on_write={}, tablet_id={}",
476
0
                _tablet_schema->keys_type(), _opts.enable_unique_key_merge_on_write,
477
0
                _tablet->tablet_id());
478
0
        DCHECK(false) << msg;
479
0
        return Status::InternalError<false>(msg);
480
0
    }
481
0
    if (_opts.rowset_ctx->partial_update_info == nullptr) {
482
0
        auto msg =
483
0
                fmt::format("partial_update_info should not be nullptr, please check, tablet_id={}",
484
0
                            _tablet->tablet_id());
485
0
        DCHECK(false) << msg;
486
0
        return Status::InternalError<false>(msg);
487
0
    }
488
0
    if (!_opts.rowset_ctx->partial_update_info->is_fixed_partial_update()) {
489
0
        auto msg = fmt::format(
490
0
                "in fixed partial update code, but update_mode={}, please check, tablet_id={}",
491
0
                _opts.rowset_ctx->partial_update_info->update_mode(), _tablet->tablet_id());
492
0
        DCHECK(false) << msg;
493
0
        return Status::InternalError<false>(msg);
494
0
    }
495
0
    if (row_pos != 0) {
496
0
        auto msg = fmt::format("row_pos should be 0, but found {}, tablet_id={}", row_pos,
497
0
                               _tablet->tablet_id());
498
0
        DCHECK(false) << msg;
499
0
        return Status::InternalError<false>(msg);
500
0
    }
501
0
    return Status::OK();
502
0
}
503
504
// for partial update, we should do following steps to fill content of block:
505
// 1. set block data to data convertor, and get all key_column's converted slice
506
// 2. get pk of input block, and read missing columns
507
//       2.1 first find key location{rowset_id, segment_id, row_id}
508
//       2.2 build read plan to read by batch
509
//       2.3 fill block
510
// 3. set columns to data convertor and then write all columns
511
Status SegmentWriter::append_block_with_partial_content(const Block* block, size_t row_pos,
512
0
                                                        size_t num_rows) {
513
0
    if (block->columns() < _tablet_schema->num_key_columns() ||
514
0
        block->columns() >= _tablet_schema->num_columns()) {
515
0
        return Status::InvalidArgument(
516
0
                fmt::format("illegal partial update block columns: {}, num key columns: {}, total "
517
0
                            "schema columns: {}",
518
0
                            block->columns(), _tablet_schema->num_key_columns(),
519
0
                            _tablet_schema->num_columns()));
520
0
    }
521
0
    RETURN_IF_ERROR(partial_update_preconditions_check(row_pos));
522
523
    // find missing column cids
524
0
    const auto& missing_cids = _opts.rowset_ctx->partial_update_info->missing_cids;
525
0
    const auto& including_cids = _opts.rowset_ctx->partial_update_info->update_cids;
526
527
    // create full block and fill with input columns
528
0
    auto full_block = _tablet_schema->create_block();
529
0
    size_t input_id = 0;
530
0
    for (auto i : including_cids) {
531
0
        full_block.replace_by_position(i, block->get_by_position(input_id++).column);
532
0
    }
533
534
0
    if (_opts.rowset_ctx->write_type != DataWriteType::TYPE_COMPACTION &&
535
0
        _tablet_schema->num_variant_columns() > 0) {
536
0
        RETURN_IF_ERROR(variant_util::parse_and_materialize_variant_columns(
537
0
                full_block, *_tablet_schema, including_cids));
538
0
    }
539
0
    RETURN_IF_ERROR(_olap_data_convertor->set_source_content_with_specifid_columns(
540
0
            &full_block, row_pos, num_rows, including_cids));
541
542
0
    bool have_input_seq_column = false;
543
    // write including columns
544
0
    std::vector<IOlapColumnDataAccessor*> key_columns;
545
0
    IOlapColumnDataAccessor* seq_column = nullptr;
546
0
    size_t segment_start_pos = 0;
547
0
    for (auto cid : including_cids) {
548
        // here we get segment column row num before append data.
549
0
        segment_start_pos = _column_writers[cid]->get_next_rowid();
550
        // olap data convertor alway start from id = 0
551
0
        auto converted_result = _olap_data_convertor->convert_column_data(cid);
552
0
        if (!converted_result.first.ok()) {
553
0
            return converted_result.first;
554
0
        }
555
0
        if (cid < _num_sort_key_columns) {
556
0
            key_columns.push_back(converted_result.second);
557
0
        } else if (_tablet_schema->has_sequence_col() &&
558
0
                   cid == _tablet_schema->sequence_col_idx()) {
559
0
            seq_column = converted_result.second;
560
0
            have_input_seq_column = true;
561
0
        }
562
0
        RETURN_IF_ERROR(_column_writers[cid]->append(converted_result.second->get_nullmap(),
563
0
                                                     converted_result.second->get_data(),
564
0
                                                     num_rows));
565
0
    }
566
567
0
    bool has_default_or_nullable = false;
568
0
    std::vector<bool> use_default_or_null_flag;
569
0
    use_default_or_null_flag.reserve(num_rows);
570
0
    const auto* delete_signs =
571
0
            BaseTablet::get_delete_sign_column_data(full_block, row_pos + num_rows);
572
573
0
    const std::vector<RowsetSharedPtr>& specified_rowsets = _mow_context->rowset_ptrs;
574
0
    std::vector<std::unique_ptr<SegmentCacheHandle>> segment_caches(specified_rowsets.size());
575
576
0
    FixedReadPlan read_plan;
577
578
    // locate rows in base data
579
0
    PartialUpdateStats stats;
580
581
0
    for (size_t block_pos = row_pos; block_pos < row_pos + num_rows; block_pos++) {
582
        // block   segment
583
        //   2   ->   0
584
        //   3   ->   1
585
        //   4   ->   2
586
        //   5   ->   3
587
        // here row_pos = 2, num_rows = 4.
588
0
        size_t delta_pos = block_pos - row_pos;
589
0
        size_t segment_pos = segment_start_pos + delta_pos;
590
0
        std::string key = _full_encode_keys(key_columns, delta_pos);
591
0
        _maybe_invalid_row_cache(key);
592
0
        if (have_input_seq_column) {
593
0
            _encode_seq_column(seq_column, delta_pos, &key);
594
0
        }
595
        // If the table have sequence column, and the include-cids don't contain the sequence
596
        // column, we need to update the primary key index builder at the end of this method.
597
        // At that time, we have a valid sequence column to encode the key with seq col.
598
0
        if (!_tablet_schema->has_sequence_col() || have_input_seq_column) {
599
0
            RETURN_IF_ERROR(_primary_key_index_builder->add_item(key));
600
0
        }
601
602
        // mark key with delete sign as deleted.
603
0
        bool have_delete_sign = (delete_signs != nullptr && delete_signs[block_pos] != 0);
604
605
0
        auto not_found_cb = [&]() {
606
0
            return _opts.rowset_ctx->partial_update_info->handle_new_key(
607
0
                    *_tablet_schema, [&]() -> std::string {
608
0
                        return block->dump_one_line(block_pos,
609
0
                                                    cast_set<int>(_num_sort_key_columns));
610
0
                    });
611
0
        };
612
0
        auto update_read_plan = [&](const RowLocation& loc) {
613
0
            read_plan.prepare_to_read(loc, segment_pos);
614
0
        };
615
0
        RETURN_IF_ERROR(probe_key_for_mow(std::move(key), segment_pos, have_input_seq_column,
616
0
                                          have_delete_sign, specified_rowsets, segment_caches,
617
0
                                          has_default_or_nullable, use_default_or_null_flag,
618
0
                                          update_read_plan, not_found_cb, stats));
619
0
    }
620
0
    CHECK_EQ(use_default_or_null_flag.size(), num_rows);
621
622
0
    if (config::enable_merge_on_write_correctness_check) {
623
0
        _tablet->add_sentinel_mark_to_delete_bitmap(_mow_context->delete_bitmap.get(),
624
0
                                                    *_mow_context->rowset_ids);
625
0
    }
626
627
    // read to fill full block
628
0
    RETURN_IF_ERROR(read_plan.fill_missing_columns(
629
0
            _opts.rowset_ctx->make_historical_row_retriever_context(), _rsid_to_rowset,
630
0
            *_tablet_schema, full_block, use_default_or_null_flag, has_default_or_nullable,
631
0
            cast_set<uint32_t>(segment_start_pos), block));
632
633
0
    if (_tablet_schema->num_variant_columns() > 0) {
634
0
        RETURN_IF_ERROR(variant_util::parse_and_materialize_variant_columns(
635
0
                full_block, *_tablet_schema, missing_cids));
636
0
    }
637
638
    // convert block to row store format
639
0
    _serialize_block_to_row_column(full_block);
640
641
    // convert missing columns and send to column writer
642
0
    RETURN_IF_ERROR(_olap_data_convertor->set_source_content_with_specifid_columns(
643
0
            &full_block, row_pos, num_rows, missing_cids));
644
0
    for (auto cid : missing_cids) {
645
0
        auto converted_result = _olap_data_convertor->convert_column_data(cid);
646
0
        if (!converted_result.first.ok()) {
647
0
            return converted_result.first;
648
0
        }
649
0
        if (_tablet_schema->has_sequence_col() && !have_input_seq_column &&
650
0
            cid == _tablet_schema->sequence_col_idx()) {
651
0
            DCHECK_EQ(seq_column, nullptr);
652
0
            seq_column = converted_result.second;
653
0
        }
654
0
        RETURN_IF_ERROR(_column_writers[cid]->append(converted_result.second->get_nullmap(),
655
0
                                                     converted_result.second->get_data(),
656
0
                                                     num_rows));
657
0
    }
658
0
    _num_rows_updated += stats.num_rows_updated;
659
0
    _num_rows_deleted += stats.num_rows_deleted;
660
0
    _num_rows_new_added += stats.num_rows_new_added;
661
0
    _num_rows_filtered += stats.num_rows_filtered;
662
0
    if (_tablet_schema->has_sequence_col() && !have_input_seq_column) {
663
0
        DCHECK_NE(seq_column, nullptr);
664
0
        if (_num_rows_written != row_pos ||
665
0
            _primary_key_index_builder->num_rows() != _num_rows_written) {
666
0
            return Status::InternalError(
667
0
                    "Correctness check failed, _num_rows_written: {}, row_pos: {}, primary key "
668
0
                    "index builder num rows: {}",
669
0
                    _num_rows_written, row_pos, _primary_key_index_builder->num_rows());
670
0
        }
671
0
        RETURN_IF_ERROR(
672
0
                _generate_primary_key_index(_key_coders, key_columns, seq_column, num_rows, false));
673
0
    }
674
675
0
    _num_rows_written += num_rows;
676
0
    DCHECK_EQ(_primary_key_index_builder->num_rows(), _num_rows_written)
677
0
            << "primary key index builder num rows(" << _primary_key_index_builder->num_rows()
678
0
            << ") not equal to segment writer's num rows written(" << _num_rows_written << ")";
679
0
    _olap_data_convertor->clear_source_content();
680
681
0
    return Status::OK();
682
0
}
683
684
3.93k
Status SegmentWriter::append_block(const Block* block, size_t row_pos, size_t num_rows) {
685
3.93k
    if (_opts.rowset_ctx->partial_update_info &&
686
3.93k
        _opts.rowset_ctx->partial_update_info->is_partial_update() &&
687
3.93k
        _opts.write_type == DataWriteType::TYPE_DIRECT &&
688
3.93k
        !_opts.rowset_ctx->is_transient_rowset_writer) {
689
0
        if (_opts.rowset_ctx->partial_update_info->is_fixed_partial_update()) {
690
0
            RETURN_IF_ERROR(append_block_with_partial_content(block, row_pos, num_rows));
691
0
        } else {
692
0
            return Status::NotSupported<false>(
693
0
                    "SegmentWriter doesn't support flexible partial update, please set "
694
0
                    "enable_vertical_segment_writer=true in be.conf on all BEs to use "
695
0
                    "VerticalSegmentWriter.");
696
0
        }
697
0
        return Status::OK();
698
0
    }
699
3.93k
    if (block->columns() < _column_writers.size()) {
700
0
        return Status::InternalError(
701
0
                "block->columns() < _column_writers.size(), block->columns()=" +
702
0
                std::to_string(block->columns()) +
703
0
                ", _column_writers.size()=" + std::to_string(_column_writers.size()) +
704
0
                ", _tablet_schema->dump_structure()=" + _tablet_schema->dump_structure());
705
0
    }
706
3.93k
    CHECK(block->columns() >= _column_writers.size())
707
0
            << ", block->columns()=" << block->columns()
708
0
            << ", _column_writers.size()=" << _column_writers.size()
709
0
            << ", _tablet_schema->dump_structure()=" << _tablet_schema->dump_structure();
710
    // Row column should be filled here when it's a directly write from memtable
711
    // or it's schema change write(since column data type maybe changed, so we should reubild)
712
3.93k
    if (_opts.write_type == DataWriteType::TYPE_DIRECT ||
713
3.93k
        _opts.write_type == DataWriteType::TYPE_SCHEMA_CHANGE) {
714
6
        _serialize_block_to_row_column(*const_cast<Block*>(block));
715
6
    }
716
717
3.93k
    if (_opts.rowset_ctx->write_type != DataWriteType::TYPE_COMPACTION &&
718
3.93k
        _tablet_schema->num_variant_columns() > 0) {
719
141
        RETURN_IF_ERROR(variant_util::parse_and_materialize_variant_columns(
720
141
                const_cast<Block&>(*block), *_tablet_schema, _column_ids));
721
141
    }
722
723
3.93k
    _olap_data_convertor->set_source_content(block, row_pos, num_rows);
724
725
    // convert column data from engine format to storage layer format
726
3.93k
    std::vector<IOlapColumnDataAccessor*> key_columns;
727
3.93k
    IOlapColumnDataAccessor* seq_column = nullptr;
728
13.5k
    for (size_t id = 0; id < _column_writers.size(); ++id) {
729
        // olap data convertor alway start from id = 0
730
9.58k
        auto converted_result = _olap_data_convertor->convert_column_data(id);
731
9.58k
        if (!converted_result.first.ok()) {
732
0
            return converted_result.first;
733
0
        }
734
9.58k
        auto cid = _column_ids[id];
735
9.58k
        if (_has_key && cid < _tablet_schema->num_key_columns()) {
736
2.55k
            key_columns.push_back(converted_result.second);
737
7.02k
        } else if (_has_key && _tablet_schema->has_sequence_col() &&
738
7.02k
                   cid == _tablet_schema->sequence_col_idx()) {
739
90
            seq_column = converted_result.second;
740
90
        }
741
9.58k
        RETURN_IF_ERROR(_column_writers[id]->append(converted_result.second->get_nullmap(),
742
9.58k
                                                    converted_result.second->get_data(), num_rows));
743
9.58k
    }
744
3.93k
    if (_opts.write_type == DataWriteType::TYPE_COMPACTION) {
745
2.21k
        RETURN_IF_ERROR(
746
2.21k
                _variant_stats_calculator->calculate_variant_stats(block, row_pos, num_rows));
747
2.21k
    }
748
749
3.93k
    RETURN_IF_ERROR(build_key_index(key_columns, seq_column, num_rows));
750
751
3.93k
    _num_rows_written += num_rows;
752
3.93k
    _olap_data_convertor->clear_source_content();
753
3.93k
    return Status::OK();
754
3.93k
}
755
756
Status SegmentWriter::build_key_index(std::vector<IOlapColumnDataAccessor*>& key_columns,
757
3.94k
                                      IOlapColumnDataAccessor* seq_column, size_t num_rows) {
758
3.94k
    if (!_has_key) {
759
981
        return Status::OK();
760
981
    }
761
762
    // find all row pos for short key indexes
763
2.95k
    std::vector<size_t> short_key_pos;
764
2.95k
    if (UNLIKELY(_short_key_row_pos == 0 && _num_rows_written == 0)) {
765
2.38k
        short_key_pos.push_back(0);
766
2.38k
    }
767
13.6k
    while (_short_key_row_pos + _opts.num_rows_per_block < _num_rows_written + num_rows) {
768
10.6k
        _short_key_row_pos += _opts.num_rows_per_block;
769
10.6k
        short_key_pos.push_back(_short_key_row_pos - _num_rows_written);
770
10.6k
    }
771
772
2.95k
    if (_is_mow_with_cluster_key()) {
773
        // For CLUSTER BY tables:
774
        // 1) generate primary key index (unique keys)
775
0
        RETURN_IF_ERROR(_generate_primary_key_index(_primary_key_coders, key_columns, seq_column,
776
0
                                                    num_rows, true));
777
        // 2) generate short key index (cluster keys)
778
0
        key_columns.clear();
779
0
        for (const auto& cid : _tablet_schema->cluster_key_uids()) {
780
0
            auto cluster_key_index = _tablet_schema->field_index(cid);
781
0
            if (cluster_key_index == -1) {
782
0
                return Status::InternalError("could not find cluster key column with unique_id=" +
783
0
                                             std::to_string(cid) + " in tablet schema");
784
0
            }
785
0
            bool found = false;
786
0
            for (auto i = 0; i < _column_ids.size(); ++i) {
787
0
                if (_column_ids[i] == cluster_key_index) {
788
0
                    auto converted_result = _olap_data_convertor->convert_column_data(i);
789
0
                    if (!converted_result.first.ok()) {
790
0
                        return converted_result.first;
791
0
                    }
792
0
                    key_columns.push_back(converted_result.second);
793
0
                    found = true;
794
0
                    break;
795
0
                }
796
0
            }
797
0
            if (!found) {
798
0
                return Status::InternalError(
799
0
                        "could not found cluster key column with unique_id=" + std::to_string(cid) +
800
0
                        ", tablet schema index=" + std::to_string(cluster_key_index));
801
0
            }
802
0
        }
803
0
        return _generate_short_key_index(key_columns, num_rows, short_key_pos);
804
0
    }
805
2.95k
    if (_is_mow()) {
806
10
        return _generate_primary_key_index(_key_coders, key_columns, seq_column, num_rows, false);
807
10
    }
808
2.94k
    return _generate_short_key_index(key_columns, num_rows, short_key_pos);
809
2.95k
}
810
811
2.07k
int64_t SegmentWriter::max_row_to_add(size_t row_avg_size_in_bytes) {
812
2.07k
    auto segment_size = estimate_segment_size();
813
2.07k
    if (segment_size >= MAX_SEGMENT_SIZE || _num_rows_written >= _opts.max_rows_per_segment)
814
354
            [[unlikely]] {
815
354
        return 0;
816
354
    }
817
1.72k
    int64_t size_rows = ((int64_t)MAX_SEGMENT_SIZE - (int64_t)segment_size) / row_avg_size_in_bytes;
818
1.72k
    int64_t count_rows = (int64_t)_opts.max_rows_per_segment - _num_rows_written;
819
820
1.72k
    return std::min(size_rows, count_rows);
821
2.07k
}
822
823
std::string SegmentWriter::_full_encode_keys(
824
6.69k
        const std::vector<IOlapColumnDataAccessor*>& key_columns, size_t pos, bool null_first) {
825
6.69k
    assert(_key_index_size.size() == _num_sort_key_columns);
826
6.69k
    assert(key_columns.size() == _num_sort_key_columns &&
827
6.69k
           _key_coders.size() == _num_sort_key_columns);
828
6.69k
    return _full_encode_keys(_key_coders, key_columns, pos, null_first);
829
6.69k
}
830
831
std::string SegmentWriter::_full_encode_keys(
832
        const std::vector<const KeyCoder*>& key_coders,
833
6.70k
        const std::vector<IOlapColumnDataAccessor*>& key_columns, size_t pos, bool null_first) {
834
6.70k
    assert(key_columns.size() == key_coders.size());
835
836
6.70k
    std::string encoded_keys;
837
6.70k
    size_t cid = 0;
838
6.70k
    for (const auto& column : key_columns) {
839
5.90k
        auto field = column->get_data_at(pos);
840
5.90k
        if (UNLIKELY(!field)) {
841
0
            if (null_first) {
842
0
                encoded_keys.push_back(KEY_NULL_FIRST_MARKER);
843
0
            } else {
844
0
                encoded_keys.push_back(KEY_NORMAL_MARKER);
845
0
            }
846
0
            ++cid;
847
0
            continue;
848
0
        }
849
5.90k
        encoded_keys.push_back(KEY_NORMAL_MARKER);
850
5.90k
        DCHECK(key_coders[cid] != nullptr);
851
5.90k
        key_coders[cid]->full_encode_ascending(field, &encoded_keys);
852
5.90k
        ++cid;
853
5.90k
    }
854
6.70k
    return encoded_keys;
855
6.70k
}
856
857
void SegmentWriter::_encode_seq_column(const IOlapColumnDataAccessor* seq_column, size_t pos,
858
0
                                       std::string* encoded_keys) {
859
0
    auto field = seq_column->get_data_at(pos);
860
    // To facilitate the use of the primary key index, encode the seq column
861
    // to the minimum value of the corresponding length when the seq column
862
    // is null
863
0
    if (UNLIKELY(!field)) {
864
0
        encoded_keys->push_back(KEY_NULL_FIRST_MARKER);
865
0
        size_t seq_col_length = _tablet_schema->column(_tablet_schema->sequence_col_idx()).length();
866
0
        encoded_keys->append(seq_col_length, KEY_MINIMAL_MARKER);
867
0
        return;
868
0
    }
869
0
    encoded_keys->push_back(KEY_NORMAL_MARKER);
870
0
    _seq_coder->full_encode_ascending(field, encoded_keys);
871
0
}
872
873
0
void SegmentWriter::_encode_rowid(const uint32_t rowid, std::string* encoded_keys) {
874
0
    encoded_keys->push_back(KEY_NORMAL_MARKER);
875
0
    _rowid_coder->full_encode_ascending(&rowid, encoded_keys);
876
0
}
877
878
std::string SegmentWriter::_encode_keys(const std::vector<IOlapColumnDataAccessor*>& key_columns,
879
13.0k
                                        size_t pos) {
880
13.0k
    assert(key_columns.size() == _num_short_key_columns);
881
882
13.0k
    std::string encoded_keys;
883
13.0k
    size_t cid = 0;
884
13.0k
    for (const auto& column : key_columns) {
885
10.0k
        auto field = column->get_data_at(pos);
886
10.0k
        if (UNLIKELY(!field)) {
887
0
            encoded_keys.push_back(KEY_NULL_FIRST_MARKER);
888
0
            ++cid;
889
0
            continue;
890
0
        }
891
10.0k
        encoded_keys.push_back(KEY_NORMAL_MARKER);
892
10.0k
        _key_coders[cid]->encode_ascending(field, _key_index_size[cid], &encoded_keys);
893
10.0k
        ++cid;
894
10.0k
    }
895
13.0k
    return encoded_keys;
896
13.0k
}
897
898
// TODO(lingbin): Currently this function does not include the size of various indexes,
899
// We should make this more precise.
900
// NOTE: This function will be called when any row of data is added, so we need to
901
// make this function efficient.
902
2.39k
uint64_t SegmentWriter::estimate_segment_size() {
903
    // footer_size(4) + checksum(4) + segment_magic(4)
904
2.39k
    uint64_t size = 12;
905
9.64k
    for (auto& column_writer : _column_writers) {
906
9.64k
        size += column_writer->estimate_buffer_size();
907
9.64k
    }
908
2.39k
    if (_is_mow_with_cluster_key()) {
909
0
        size += _primary_key_index_builder->size() + _short_key_index_builder->size();
910
2.39k
    } else if (_is_mow()) {
911
0
        size += _primary_key_index_builder->size();
912
2.39k
    } else {
913
2.39k
        size += _short_key_index_builder->size();
914
2.39k
    }
915
916
    // update the mem_tracker of segment size
917
2.39k
    _mem_tracker->consume(size - _mem_tracker->consumption());
918
2.39k
    return size;
919
2.39k
}
920
921
3.37k
Status SegmentWriter::finalize_columns_data() {
922
3.37k
    if (_has_key) {
923
2.44k
        _row_count = _num_rows_written;
924
2.44k
    } else {
925
924
        DCHECK(_row_count == _num_rows_written)
926
0
                << "_row_count != _num_rows_written:" << _row_count << " vs. " << _num_rows_written;
927
924
        if (_row_count != _num_rows_written) {
928
0
            std::stringstream ss;
929
0
            ss << "_row_count != _num_rows_written:" << _row_count << " vs. " << _num_rows_written;
930
0
            LOG(WARNING) << ss.str();
931
0
            return Status::InternalError(ss.str());
932
0
        }
933
924
    }
934
3.37k
    _num_rows_written = 0;
935
936
10.2k
    for (auto& column_writer : _column_writers) {
937
10.2k
        RETURN_IF_ERROR(column_writer->finish());
938
10.2k
    }
939
3.37k
    RETURN_IF_ERROR(_write_data());
940
941
3.37k
    return Status::OK();
942
3.37k
}
943
944
3.37k
Status SegmentWriter::finalize_columns_index(uint64_t* index_size) {
945
3.37k
    uint64_t index_start = _file_writer->bytes_appended();
946
3.37k
    RETURN_IF_ERROR(_write_ordinal_index());
947
3.37k
    RETURN_IF_ERROR(_write_zone_map());
948
3.37k
    RETURN_IF_ERROR(_write_inverted_index());
949
3.37k
    RETURN_IF_ERROR(_write_ann_index());
950
3.37k
    RETURN_IF_ERROR(_write_bloom_filter_index());
951
952
3.37k
    *index_size = _file_writer->bytes_appended() - index_start;
953
3.37k
    if (_has_key) {
954
2.44k
        if (_is_mow_with_cluster_key()) {
955
            // 1. sort primary keys
956
0
            std::sort(_primary_keys.begin(), _primary_keys.end());
957
            // 2. write primary keys index
958
0
            std::string last_key;
959
0
            for (const auto& key : _primary_keys) {
960
0
                DCHECK(key.compare(last_key) > 0)
961
0
                        << "found duplicate key or key is not sorted! current key: " << key
962
0
                        << ", last key: " << last_key;
963
0
                RETURN_IF_ERROR(_primary_key_index_builder->add_item(key));
964
0
                last_key = key;
965
0
            }
966
967
0
            RETURN_IF_ERROR(_write_short_key_index());
968
0
            *index_size = _file_writer->bytes_appended() - index_start;
969
0
            RETURN_IF_ERROR(_write_primary_key_index());
970
0
            *index_size += _primary_key_index_builder->disk_size();
971
2.44k
        } else if (_is_mow()) {
972
69
            RETURN_IF_ERROR(_write_primary_key_index());
973
            // IndexedColumnWriter write data pages mixed with segment data, we should use
974
            // the stat from primary key index builder.
975
69
            *index_size += _primary_key_index_builder->disk_size();
976
2.37k
        } else {
977
2.37k
            RETURN_IF_ERROR(_write_short_key_index());
978
2.37k
            *index_size = _file_writer->bytes_appended() - index_start;
979
2.37k
        }
980
2.44k
    }
981
    // reset all column writers and data_conveter
982
3.37k
    clear();
983
984
3.37k
    return Status::OK();
985
3.37k
}
986
987
2.44k
Status SegmentWriter::finalize_footer(uint64_t* segment_file_size) {
988
2.44k
    RETURN_IF_ERROR(_write_footer());
989
    // finish
990
2.44k
    RETURN_IF_ERROR(_file_writer->close(true));
991
2.44k
    *segment_file_size = _file_writer->bytes_appended();
992
2.44k
    if (*segment_file_size == 0) {
993
0
        return Status::Corruption("Bad segment, file size = 0");
994
0
    }
995
2.44k
    return Status::OK();
996
2.44k
}
997
998
1.25k
Status SegmentWriter::finalize(uint64_t* segment_file_size, uint64_t* index_size) {
999
1.25k
    MonotonicStopWatch timer;
1000
1.25k
    timer.start();
1001
    // check disk capacity
1002
1.25k
    if (_data_dir != nullptr && _data_dir->reach_capacity_limit((int64_t)estimate_segment_size())) {
1003
0
        return Status::Error<DISK_REACH_CAPACITY_LIMIT>("disk {} exceed capacity limit, path: {}",
1004
0
                                                        _data_dir->path_hash(), _data_dir->path());
1005
0
    }
1006
    // write data
1007
1.25k
    RETURN_IF_ERROR(finalize_columns_data());
1008
    // Get the index start before finalize_footer since this function would write new data.
1009
1.25k
    uint64_t index_start = _file_writer->bytes_appended();
1010
    // write index
1011
1.25k
    RETURN_IF_ERROR(finalize_columns_index(index_size));
1012
    // write footer
1013
1.25k
    RETURN_IF_ERROR(finalize_footer(segment_file_size));
1014
1015
1.25k
    if (timer.elapsed_time() > 5000000000l) {
1016
0
        LOG(INFO) << "segment flush consumes a lot time_ns " << timer.elapsed_time()
1017
0
                  << ", segmemt_size " << *segment_file_size;
1018
0
    }
1019
    // When the cache type is not ttl(expiration time == 0), the data should be split into normal cache queue
1020
    // and index cache queue
1021
1.25k
    if (auto* cache_builder = _file_writer->cache_builder(); cache_builder != nullptr &&
1022
1.25k
                                                             cache_builder->_expiration_time == 0 &&
1023
1.25k
                                                             config::is_cloud_mode()) {
1024
0
        auto size = *index_size + *segment_file_size;
1025
0
        auto holder = cache_builder->allocate_cache_holder(index_start, size, _tablet->tablet_id());
1026
0
        for (auto& segment : holder->file_blocks) {
1027
0
            static_cast<void>(segment->change_cache_type(io::FileCacheType::INDEX));
1028
0
        }
1029
0
    }
1030
1.25k
    return Status::OK();
1031
1.25k
}
1032
1033
3.39k
void SegmentWriter::clear() {
1034
10.2k
    for (auto& column_writer : _column_writers) {
1035
10.2k
        column_writer.reset();
1036
10.2k
    }
1037
3.39k
    _column_writers.clear();
1038
3.39k
    _column_ids.clear();
1039
3.39k
    _olap_data_convertor.reset();
1040
3.39k
}
1041
1042
// write column data to file one by one
1043
3.37k
Status SegmentWriter::_write_data() {
1044
10.2k
    for (auto& column_writer : _column_writers) {
1045
10.2k
        RETURN_IF_ERROR(column_writer->write_data());
1046
1047
10.2k
        auto* column_meta = column_writer->get_column_meta();
1048
10.2k
        DCHECK(column_meta != nullptr);
1049
10.2k
        column_meta->set_compressed_data_bytes(
1050
10.2k
                (column_meta->has_compressed_data_bytes() ? column_meta->compressed_data_bytes()
1051
10.2k
                                                          : 0) +
1052
10.2k
                column_writer->get_total_compressed_data_pages_bytes());
1053
10.2k
        column_meta->set_uncompressed_data_bytes(
1054
10.2k
                (column_meta->has_uncompressed_data_bytes() ? column_meta->uncompressed_data_bytes()
1055
10.2k
                                                            : 0) +
1056
10.2k
                column_writer->get_total_uncompressed_data_pages_bytes());
1057
10.2k
        column_meta->set_raw_data_bytes(
1058
10.2k
                (column_meta->has_raw_data_bytes() ? column_meta->raw_data_bytes() : 0) +
1059
10.2k
                column_writer->get_raw_data_bytes());
1060
10.2k
    }
1061
3.37k
    return Status::OK();
1062
3.37k
}
1063
1064
// write ordinal index after data has been written
1065
3.37k
Status SegmentWriter::_write_ordinal_index() {
1066
10.2k
    for (auto& column_writer : _column_writers) {
1067
10.2k
        RETURN_IF_ERROR(column_writer->write_ordinal_index());
1068
10.2k
    }
1069
3.37k
    return Status::OK();
1070
3.37k
}
1071
1072
3.37k
Status SegmentWriter::_write_zone_map() {
1073
10.2k
    for (auto& column_writer : _column_writers) {
1074
10.2k
        RETURN_IF_ERROR(column_writer->write_zone_map());
1075
10.2k
    }
1076
3.37k
    return Status::OK();
1077
3.37k
}
1078
1079
3.37k
Status SegmentWriter::_write_inverted_index() {
1080
10.2k
    for (auto& column_writer : _column_writers) {
1081
10.2k
        RETURN_IF_ERROR(column_writer->write_inverted_index());
1082
10.2k
    }
1083
3.37k
    return Status::OK();
1084
3.37k
}
1085
1086
3.37k
Status SegmentWriter::_write_ann_index() {
1087
10.2k
    for (auto& column_writer : _column_writers) {
1088
10.2k
        RETURN_IF_ERROR(column_writer->write_ann_index());
1089
10.2k
    }
1090
3.37k
    return Status::OK();
1091
3.37k
}
1092
1093
3.37k
Status SegmentWriter::_write_bloom_filter_index() {
1094
10.2k
    for (auto& column_writer : _column_writers) {
1095
10.2k
        RETURN_IF_ERROR(column_writer->write_bloom_filter_index());
1096
10.2k
    }
1097
3.37k
    return Status::OK();
1098
3.37k
}
1099
1100
2.37k
Status SegmentWriter::_write_short_key_index() {
1101
2.37k
    std::vector<Slice> body;
1102
2.37k
    PageFooterPB footer;
1103
2.37k
    RETURN_IF_ERROR(_short_key_index_builder->finalize(_row_count, &body, &footer));
1104
2.37k
    PagePointer pp;
1105
    // short key index page is not compressed right now
1106
2.37k
    RETURN_IF_ERROR(PageIO::write_page(_file_writer, body, footer, &pp));
1107
2.37k
    pp.to_proto(_footer.mutable_short_key_index_page());
1108
2.37k
    return Status::OK();
1109
2.37k
}
1110
1111
69
Status SegmentWriter::_write_primary_key_index() {
1112
69
    CHECK_EQ(_primary_key_index_builder->num_rows(), _row_count);
1113
69
    return _primary_key_index_builder->finalize(_footer.mutable_primary_key_index_meta());
1114
69
}
1115
1116
2.44k
Status SegmentWriter::_write_footer() {
1117
2.44k
    _footer.set_num_rows(_row_count);
1118
    // Decide whether to externalize ColumnMetaPB by tablet default, and stamp footer version
1119
2.44k
    if (_tablet_schema->storage_format() == TabletStorageFormatPB::TABLET_STORAGE_FORMAT_V3) {
1120
72
        _footer.set_version(SEGMENT_FOOTER_VERSION_V3_EXT_COL_META);
1121
72
        VLOG_DEBUG << "use external column meta";
1122
        // External ColumnMetaPB writing (optional)
1123
72
        RETURN_IF_ERROR(ExternalColMetaUtil::write_external_column_meta(
1124
72
                _file_writer, &_footer, _opts.compression_type,
1125
72
                [this](const std::vector<Slice>& slices) { return _write_raw_data(slices); }));
1126
72
    }
1127
1128
    // Footer := SegmentFooterPB, FooterPBSize(4), FooterPBChecksum(4), MagicNumber(4)
1129
2.44k
    std::string footer_buf;
1130
2.44k
    VLOG_DEBUG << "footer " << _footer.DebugString();
1131
2.44k
    if (!_footer.SerializeToString(&footer_buf)) {
1132
0
        return Status::InternalError("failed to serialize segment footer");
1133
0
    }
1134
1135
2.44k
    faststring fixed_buf;
1136
    // footer's size
1137
2.44k
    put_fixed32_le(&fixed_buf, cast_set<uint32_t>(footer_buf.size()));
1138
    // footer's checksum
1139
2.44k
    uint32_t checksum = crc32c::Crc32c(footer_buf.data(), footer_buf.size());
1140
2.44k
    put_fixed32_le(&fixed_buf, checksum);
1141
    // Append magic number. we don't write magic number in the header because
1142
    // that will need an extra seek when reading
1143
2.44k
    fixed_buf.append(k_segment_magic, k_segment_magic_length);
1144
1145
2.44k
    std::vector<Slice> slices {footer_buf, fixed_buf};
1146
2.44k
    return _write_raw_data(slices);
1147
2.44k
}
1148
1149
4.17k
Status SegmentWriter::_write_raw_data(const std::vector<Slice>& slices) {
1150
4.17k
    RETURN_IF_ERROR(_file_writer->appendv(&slices[0], slices.size()));
1151
4.17k
    return Status::OK();
1152
4.17k
}
1153
1154
2.44k
Slice SegmentWriter::min_encoded_key() {
1155
2.44k
    return (_primary_key_index_builder == nullptr) ? Slice(_min_key.data(), _min_key.size())
1156
2.44k
                                                   : _primary_key_index_builder->min_key();
1157
2.44k
}
1158
2.44k
Slice SegmentWriter::max_encoded_key() {
1159
2.44k
    return (_primary_key_index_builder == nullptr) ? Slice(_max_key.data(), _max_key.size())
1160
2.44k
                                                   : _primary_key_index_builder->max_key();
1161
2.44k
}
1162
1163
211
void SegmentWriter::set_min_max_key(const Slice& key) {
1164
211
    if (UNLIKELY(_is_first_row)) {
1165
6
        _min_key.append(key.get_data(), key.get_size());
1166
6
        _is_first_row = false;
1167
6
    }
1168
211
    if (key.compare(_max_key) > 0) {
1169
211
        _max_key.clear();
1170
211
        _max_key.append(key.get_data(), key.get_size());
1171
211
    }
1172
211
}
1173
1174
2.94k
void SegmentWriter::set_min_key(const Slice& key) {
1175
2.94k
    if (UNLIKELY(_is_first_row)) {
1176
2.37k
        _min_key.append(key.get_data(), key.get_size());
1177
2.37k
        _is_first_row = false;
1178
2.37k
    }
1179
2.94k
}
1180
1181
2.94k
void SegmentWriter::set_max_key(const Slice& key) {
1182
2.94k
    _max_key.clear();
1183
2.94k
    _max_key.append(key.get_data(), key.get_size());
1184
2.94k
}
1185
1186
0
void SegmentWriter::set_mow_context(std::shared_ptr<MowContext> mow_context) {
1187
0
    _mow_context = mow_context;
1188
0
}
1189
1190
Status SegmentWriter::_generate_primary_key_index(
1191
        const std::vector<const KeyCoder*>& primary_key_coders,
1192
        const std::vector<IOlapColumnDataAccessor*>& primary_key_columns,
1193
10
        IOlapColumnDataAccessor* seq_column, size_t num_rows, bool need_sort) {
1194
10
    if (!need_sort) { // mow table without cluster key
1195
10
        std::string last_key;
1196
810
        for (size_t pos = 0; pos < num_rows; pos++) {
1197
            // use _key_coders
1198
800
            std::string key = _full_encode_keys(primary_key_columns, pos);
1199
800
            _maybe_invalid_row_cache(key);
1200
800
            if (_tablet_schema->has_sequence_col()) {
1201
0
                _encode_seq_column(seq_column, pos, &key);
1202
0
            }
1203
800
            DCHECK(key.compare(last_key) > 0)
1204
0
                    << "found duplicate key or key is not sorted! current key: " << key
1205
0
                    << ", last key: " << last_key;
1206
800
            RETURN_IF_ERROR(_primary_key_index_builder->add_item(key));
1207
800
            last_key = std::move(key);
1208
800
        }
1209
10
    } else { // mow table with cluster key
1210
        // generate primary keys in memory
1211
0
        for (uint32_t pos = 0; pos < num_rows; pos++) {
1212
0
            std::string key = _full_encode_keys(primary_key_coders, primary_key_columns, pos);
1213
0
            _maybe_invalid_row_cache(key);
1214
0
            if (_tablet_schema->has_sequence_col()) {
1215
0
                _encode_seq_column(seq_column, pos, &key);
1216
0
            }
1217
0
            _encode_rowid(pos + _num_rows_written, &key);
1218
0
            _primary_keys_size += key.size();
1219
0
            _primary_keys.emplace_back(std::move(key));
1220
0
        }
1221
0
    }
1222
10
    return Status::OK();
1223
10
}
1224
1225
Status SegmentWriter::_generate_short_key_index(std::vector<IOlapColumnDataAccessor*>& key_columns,
1226
                                                size_t num_rows,
1227
2.94k
                                                const std::vector<size_t>& short_key_pos) {
1228
    // use _key_coders
1229
2.94k
    set_min_key(_full_encode_keys(key_columns, 0));
1230
2.94k
    set_max_key(_full_encode_keys(key_columns, num_rows - 1));
1231
2.94k
    DCHECK(Slice(_max_key.data(), _max_key.size())
1232
0
                   .compare(Slice(_min_key.data(), _min_key.size())) >= 0)
1233
0
            << "key is not sorted! min key: " << _min_key << ", max key: " << _max_key;
1234
1235
2.94k
    key_columns.resize(_num_short_key_columns);
1236
2.94k
    std::string last_key;
1237
13.0k
    for (const auto pos : short_key_pos) {
1238
13.0k
        std::string key = _encode_keys(key_columns, pos);
1239
13.0k
        DCHECK(key.compare(last_key) >= 0)
1240
0
                << "key is not sorted! current key: " << key << ", last key: " << last_key;
1241
13.0k
        RETURN_IF_ERROR(_short_key_index_builder->add_item(key));
1242
13.0k
        last_key = std::move(key);
1243
13.0k
    }
1244
2.94k
    return Status::OK();
1245
2.94k
}
1246
1247
281k
inline bool SegmentWriter::_is_mow() {
1248
281k
    return _tablet_schema->keys_type() == UNIQUE_KEYS && _opts.enable_unique_key_merge_on_write;
1249
281k
}
1250
1251
139k
inline bool SegmentWriter::_is_mow_with_cluster_key() {
1252
139k
    return _is_mow() && !_tablet_schema->cluster_key_uids().empty();
1253
139k
}
1254
1255
} // namespace segment_v2
1256
} // namespace doris