[refactor](storage) extract RowKeyEncoder: dedup key encoding across both segment writers

be/src/storage/key/row_key_encoder.cpp

@@ -0,0 +1,137 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "storage/key/row_key_encoder.h"
+
+#include <cassert>
+
+#include "common/cast_set.h"
+#include "common/compiler_util.h" // IWYU pragma: keep
+#include "common/consts.h"
+#include "common/logging.h"
+#include "storage/iterator/olap_data_convertor.h"
+#include "storage/key_coder.h"
+#include "storage/tablet/tablet_schema.h"
+
+namespace doris {
+
+RowKeyEncoder::RowKeyEncoder(const TabletSchema& schema, bool mow) {
+    for (size_t cid = 0; cid < schema.num_key_columns(); ++cid) {
+        _primary_key_coders.push_back(get_key_coder(schema.column(cid).type()));
+    }
+    _num_short_key_columns = schema.num_short_key_columns();
+    // encode the sequence id into the primary key index
+    if (schema.has_sequence_col()) {
+        const auto& column = schema.column(schema.sequence_col_idx());
+        _seq_coder = get_key_coder(column.type());
+        _seq_col_length = column.length();
+    }
+    if (mow && !schema.cluster_key_uids().empty()) {
+        _rowid_coder = get_key_coder(FieldType::OLAP_FIELD_TYPE_UNSIGNED_INT);
+        for (auto uid : schema.cluster_key_uids()) {
+            _add_sort_key_column(schema.column_by_uid(uid));
+        }
+    } else {
+        for (size_t cid = 0; cid < schema.num_key_columns(); ++cid) {
+            _add_sort_key_column(schema.column(cid));
+        }
+    }
+}
+
+void RowKeyEncoder::_add_sort_key_column(const TabletColumn& column) {
+    _sort_key_coders.push_back(get_key_coder(column.type()));
+    _sort_key_index_size.push_back(cast_set<uint16_t>(column.index_length()));
+}
+
+std::string RowKeyEncoder::full_encode(const std::vector<IOlapColumnDataAccessor*>& key_columns,
+                                       size_t pos) const {
+    assert(_sort_key_index_size.size() == _sort_key_coders.size());
+    assert(key_columns.size() == _sort_key_coders.size());
+    return _full_encode(_sort_key_coders, key_columns, pos);
+}
+
+std::string RowKeyEncoder::full_encode_primary_keys(
+        const std::vector<IOlapColumnDataAccessor*>& key_columns, size_t pos) const {
+    return _full_encode(_primary_key_coders, key_columns, pos);
+}
+
+namespace {
+// Shared row-key encoding base: for each key column, write a null marker for
+// a null value, otherwise a normal marker followed by whatever `encode_field`
+// appends. `encode_field(cid, field, out)` is the only thing that differs between
+// the full key encode and the short-key prefix encode.
+template <typename EncodeField>
+std::string encode_key_columns(const std::vector<IOlapColumnDataAccessor*>& key_columns, size_t pos,
+                               EncodeField&& encode_field) {
+    std::string encoded_keys;
+    size_t cid = 0;
+    for (const auto& column : key_columns) {
+        const auto* field = column->get_data_at(pos);
+        if (UNLIKELY(!field)) {
+            encoded_keys.push_back(KeyConsts::KEY_NULL_FIRST_MARKER);
+            ++cid;
+            continue;
+        }
+        encoded_keys.push_back(KeyConsts::KEY_NORMAL_MARKER);
+        encode_field(cid, field, &encoded_keys);
+        ++cid;
+    }
+    return encoded_keys;
+}
+} // namespace
+
+std::string RowKeyEncoder::_full_encode(const std::vector<const KeyCoder*>& key_coders,
+                                        const std::vector<IOlapColumnDataAccessor*>& key_columns,
+                                        size_t pos) {
+    assert(key_columns.size() == key_coders.size());
+    return encode_key_columns(key_columns, pos,
+                              [&](size_t cid, const void* field, std::string* out) {
+                                  DCHECK(key_coders[cid] != nullptr);
+                                  key_coders[cid]->full_encode_ascending(field, out);
+                              });
+}
+
+std::string RowKeyEncoder::encode_short_keys(
+        const std::vector<IOlapColumnDataAccessor*>& key_columns, size_t pos) const {
+    assert(key_columns.size() == _num_short_key_columns);
+    assert(key_columns.size() <= _sort_key_coders.size());
+    return encode_key_columns(
+            key_columns, pos, [&](size_t cid, const void* field, std::string* out) {
+                _sort_key_coders[cid]->encode_ascending(field, _sort_key_index_size[cid], out);
+            });
+}
+
+void RowKeyEncoder::append_seq_suffix(std::string* encoded_keys,
+                                      const IOlapColumnDataAccessor* seq_column, size_t pos) const {
+    const auto* field = seq_column->get_data_at(pos);
+    // So the primary key index can still use it, encode a null seq column as
+    // the smallest value of its length.
+    if (UNLIKELY(!field)) {
+        encoded_keys->push_back(KeyConsts::KEY_NULL_FIRST_MARKER);
+        encoded_keys->append(_seq_col_length, KeyConsts::KEY_MINIMAL_MARKER);
+        return;
+    }
+    encoded_keys->push_back(KeyConsts::KEY_NORMAL_MARKER);
+    _seq_coder->full_encode_ascending(field, encoded_keys);
+}
+
+void RowKeyEncoder::append_rowid_suffix(std::string* encoded_keys, uint32_t rowid) const {
+    encoded_keys->push_back(KeyConsts::KEY_NORMAL_MARKER);
+    _rowid_coder->full_encode_ascending(&rowid, encoded_keys);
+}
+
+} // namespace doris

be/src/storage/key/row_key_encoder.h

@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace doris {
+
+class IOlapColumnDataAccessor;
+class KeyCoder;
+class TabletColumn;
+class TabletSchema;
+
+// Encodes rows into the sortable binary key format shared by the short key
+// index and the primary key index. Each column is encoded as a one byte
+// marker (KeyConsts) followed by the KeyCoder encoded value, a null column
+// is encoded as KEY_NULL_FIRST_MARKER without value bytes.
+class RowKeyEncoder {
+public:
+    RowKeyEncoder(const TabletSchema& schema, bool mow);
+
+    // Encode the sort key columns at `pos` with full length.
+    std::string full_encode(const std::vector<IOlapColumnDataAccessor*>& key_columns,
+                            size_t pos) const;
+
+    // Encode the primary key columns at `pos` with full length, producing the
+    // key stored in and probed against the primary key index.
+    std::string full_encode_primary_keys(const std::vector<IOlapColumnDataAccessor*>& key_columns,
+                                         size_t pos) const;
+
+    // Encode the short key columns at `pos`, each column truncated to its
+    // index length.
+    std::string encode_short_keys(const std::vector<IOlapColumnDataAccessor*>& key_columns,
+                                  size_t pos) const;
+
+    // Append the encoded sequence column at `pos` to `encoded_keys`. A null
+    // sequence value is encoded as the minimal value of the column length so
+    // that it sorts first in the primary key index.
+    void append_seq_suffix(std::string* encoded_keys, const IOlapColumnDataAccessor* seq_column,
+                           size_t pos) const;
+
+    // Append the encoded row id to `encoded_keys`, only used by mow tables
+    // with cluster keys.
+    void append_rowid_suffix(std::string* encoded_keys, uint32_t rowid) const;
+
+    size_t num_sort_key_columns() const { return _sort_key_coders.size(); }
+
+private:
+    static std::string _full_encode(const std::vector<const KeyCoder*>& key_coders,
+                                    const std::vector<IOlapColumnDataAccessor*>& key_columns,
+                                    size_t pos);
+
+    void _add_sort_key_column(const TabletColumn& column);
+
+    // The sort-key view: whatever the segment sorts by. Cluster key columns
+    // for mow tables with cluster keys, primary key columns otherwise. Used by
+    // full_encode() and encode_short_keys().
+    std::vector<const KeyCoder*> _sort_key_coders;
+    std::vector<uint16_t> _sort_key_index_size;
+    // The primary-key view: always the primary key columns. The primary key
+    // index is built on these even when the segment sorts by cluster keys;
+    // used by full_encode_primary_keys().
+    std::vector<const KeyCoder*> _primary_key_coders;
+    const KeyCoder* _seq_coder = nullptr;
+    const KeyCoder* _rowid_coder = nullptr;
+    size_t _seq_col_length = 0;
+    size_t _num_short_key_columns = 0;
+};
+
+} // namespace doris

be/src/storage/partial_update_info.cpp

@@ -29,6 +29,7 @@
 #include "core/data_type/data_type_number.h" // IWYU pragma: keep
 #include "core/value/bitmap_value.h"
 #include "storage/iterator/olap_data_convertor.h"
+#include "storage/key/row_key_encoder.h"
 #include "storage/olap_common.h"
 #include "storage/rowset/rowset.h"
 #include "storage/rowset/rowset_writer_context.h"
@@ -915,7 +916,7 @@ Status BlockAggregator::aggregate_rows(
             // Discard all the rows whose seq value is smaller than previous_encoded_seq_value.
             if (row_has_sequence_col) {
                 std::string seq_val {};
-                _writer._encode_seq_column(seq_column, pos, &seq_val);
+                _writer._key_encoder.append_seq_suffix(&seq_val, seq_column, pos);
                 if (Slice {seq_val}.compare(Slice {previous_encoded_seq_value}) < 0) {
                     continue;
                 }
@@ -932,7 +933,7 @@ Status BlockAggregator::aggregate_rows(
         if (row_has_sequence_col) {
             std::string seq_val {};
             // for rows that don't specify seqeunce col, seq_val will be encoded to minial value
-            _writer._encode_seq_column(seq_column, pos, &seq_val);
+            _writer._key_encoder.append_seq_suffix(&seq_val, seq_column, pos);
             cur_seq_val = std::move(seq_val);
         } else {
             cur_seq_val.clear();
@@ -950,7 +951,7 @@ Status BlockAggregator::aggregate_rows(
             append_or_merge_row(output_block, block, rid, skip_bitmap, have_delete_sign);
         } else {
             std::string seq_val {};
-            _writer._encode_seq_column(seq_column, rid, &seq_val);
+            _writer._key_encoder.append_seq_suffix(&seq_val, seq_column, rid);
             if (Slice {seq_val}.compare(Slice {cur_seq_val}) >= 0) {
                 append_or_merge_row(output_block, block, rid, skip_bitmap, have_delete_sign);
                 cur_seq_val = std::move(seq_val);
@@ -981,7 +982,7 @@ Status BlockAggregator::aggregate_for_sequence_column(
     int same_key_rows {0};
     std::string previous_key {};
     for (int block_pos {0}; block_pos < num_rows; block_pos++) {
-        std::string key = _writer._full_encode_keys(key_columns, block_pos);
+        std::string key = _writer._key_encoder.full_encode(key_columns, block_pos);
         if (block_pos > 0 && previous_key == key) {
             same_key_rows++;
         } else {
@@ -1061,7 +1062,7 @@ Status BlockAggregator::aggregate_for_insert_after_delete(
     for (size_t block_pos {0}; block_pos < num_rows; block_pos++) {
         size_t delta_pos = block_pos;
         auto& skip_bitmap = skip_bitmaps->at(block_pos);
-        std::string key = _writer._full_encode_keys(key_columns, delta_pos);
+        std::string key = _writer._key_encoder.full_encode(key_columns, delta_pos);
         bool have_delete_sign =
                 (!skip_bitmap.contains(delete_sign_col_unique_id) && delete_signs[block_pos] != 0);
         if (delta_pos > 0 && previous_key == key) {