[Enhancement] Expose Kafka/Pulsar message metadata via an INCLUDE METADATA clause in Routine Load (#73840)

Signed-off-by: Evgeniy Shishkin <eshishki@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 0796ea6)

Author: 2 people

be/src/exec/CMakeLists.txt

@@ -129,6 +129,7 @@ set(EXEC_FILES
     file_scanner/json_scanner.cpp
     file_scanner/orc_scanner.cpp
     file_scanner/parquet_scanner.cpp
+    file_scanner/stream_source_meta.cpp
     schema_scanner/schema_analyze_status.cpp
     schema_scanner/schema_tables_scanner.cpp
     schema_scanner/schema_dummy_scanner.cpp

be/src/exec/file_scanner/avro_scanner.cpp

@@ -26,6 +26,7 @@
 #include "column/chunk.h"
 #include "column/column_helper.h"
 #include "exec/file_scanner/json_scanner.h"
+#include "exec/file_scanner/stream_source_meta.h"
 #include "exec/json_parser.h"
 #include "exprs/cast_expr.h"
 #include "exprs/column_ref.h"
@@ -160,6 +161,8 @@ Status AvroScanner::open() {
     }
     ++_counter->num_files_read;
 
+    _meta_col_by_slot_id = build_stream_source_meta_columns(_scan_range.params.stream_source_meta_columns);
+    int column_index = 0;
     for (size_t i = 0; i < _src_slot_descriptors.size(); ++i) {
         const auto& desc = _src_slot_descriptors[i];
         if (desc == nullptr) {
@@ -169,6 +172,13 @@ Status AvroScanner::open() {
         // Remember the intermediate avro type per slot so the no-jsonpath path can dispatch
         // _construct_column on the source column's type rather than the destination type.
         _slot_id_to_avro_type.emplace(desc->id(), _avro_types[i]);
+        // A hidden metadata slot is filled from the message meta, not the payload. The by-name null-fill
+        // path iterates chunk columns, so key its descriptor by chunk column index (the append order in
+        // _create_src_chunk, which skips null slots just like this loop).
+        if (auto m = _meta_col_by_slot_id.find(desc->id()); m != _meta_col_by_slot_id.end()) {
+            _meta_col_by_index.emplace(column_index, m->second);
+        }
+        column_index++;
     }
     _init_data_idx_to_slot_once = false;
     return Status::OK();
@@ -205,20 +215,33 @@ void AvroScanner::_report_error(const std::string& line, const std::string& err_
     _state->append_error_msg_to_file(line, err_msg);
 }
 
-Status AvroScanner::_construct_row(const avro_value_t& avro_value, Chunk* chunk) {
+Status AvroScanner::_construct_row(const avro_value_t& avro_value, Chunk* chunk, const StreamMessageMeta* meta) {
     size_t slot_size = _src_slot_descriptors.size();
     size_t jsonpath_size = _json_paths.size();
+    // Hidden metadata columns carry no jsonpath, so they are transparent to the positional
+    // slot->jsonpath mapping: a jsonpath maps to the i-th non-metadata slot. This keeps payload columns
+    // aligned even when a metadata column sits before jsonpath-backed columns.
+    size_t meta_count = 0;
     for (size_t i = 0; i < slot_size; i++) {
         if (_src_slot_descriptors[i] == nullptr) {
             continue;
         }
         auto column = down_cast<NullableColumn*>(chunk->get_column_raw_ptr_by_slot_id(_src_slot_descriptors[i]->id()));
-        if (UNLIKELY(i >= jsonpath_size)) {
+        // Hidden source-metadata slots are filled from the message meta (by slot id), not the payload,
+        // before the jsonpath extraction below.
+        if (auto it = _meta_col_by_slot_id.find(_src_slot_descriptors[i]->id());
+            UNLIKELY(it != _meta_col_by_slot_id.end())) {
+            RETURN_IF_ERROR(fill_stream_source_meta_column(it->second.kind, meta, column));
+            meta_count++;
+            continue;
+        }
+        size_t jp = i - meta_count;
+        if (UNLIKELY(jp >= jsonpath_size)) {
             column->append_nulls(1);
             continue;
         }
         avro_value_t output_value;
-        auto st = _extract_field(avro_value, _json_paths[i], &output_value);
+        auto st = _extract_field(avro_value, _json_paths[jp], &output_value);
         if (LIKELY(st.ok())) {
             // Dispatch on the intermediate avro type (which matches the source column created in
             // _create_src_chunk), not the destination type. Otherwise a complex column whose
@@ -241,6 +264,9 @@ Status AvroScanner::_parse_avro(Chunk* chunk, const std::shared_ptr<SequentialFi
     DCHECK_EQ(0, chunk->num_rows());
     for (size_t num_rows = chunk->num_rows(); num_rows < capacity; /**/) {
         avro_value_t avro_value;
+        // Routine-load source-metadata for this message (null for non-routine-load); read from the pipe
+        // buffer below, or injected by the test in BE_TEST.
+        const StreamMessageMeta* meta = nullptr;
 #ifdef BE_TEST
         // In general, we want to test component injection schemastr.
         avro_schema_error_t error;
@@ -273,6 +299,7 @@ Status AvroScanner::_parse_avro(Chunk* chunk, const std::shared_ptr<SequentialFi
             return Status::InternalError(err_msg);
         }
         free(avro_as_json);
+        meta = _test_meta;
 #else
         const uint8_t* data{};
         size_t length = 0;
@@ -284,6 +311,7 @@ Status AvroScanner::_parse_avro(Chunk* chunk, const std::shared_ptr<SequentialFi
         }
         data = reinterpret_cast<uint8_t*>(_parser_buf->ptr);
         length = _parser_buf->remaining();
+        meta = stream_source_meta_of(_parser_buf);
         serdes_schema_t* schema;
         serdes_err_t err =
                 serdes_deserialize_avro(_serdes, &avro_value, &schema, data, length, _err_buf, sizeof(_err_buf));
@@ -299,7 +327,7 @@ Status AvroScanner::_parse_avro(Chunk* chunk, const std::shared_ptr<SequentialFi
         size_t chunk_row_num = chunk->num_rows();
         Status st = Status::OK();
         if (!_json_paths.empty()) {
-            st = _construct_row(avro_value, chunk);
+            st = _construct_row(avro_value, chunk, meta);
         } else {
             if (!_init_data_idx_to_slot_once) {
                 size_t element_count;
@@ -320,7 +348,7 @@ Status AvroScanner::_parse_avro(Chunk* chunk, const std::shared_ptr<SequentialFi
 
                 _init_data_idx_to_slot_once = true;
             }
-            st = _construct_row_without_jsonpath(avro_value, chunk);
+            st = _construct_row_without_jsonpath(avro_value, chunk, meta);
         }
         if (!st.ok()) {
             if (_counter->num_rows_filtered++ < MAX_ERROR_LINES_IN_FILE) {
@@ -338,7 +366,8 @@ Status AvroScanner::_parse_avro(Chunk* chunk, const std::shared_ptr<SequentialFi
     return Status::OK();
 }
 
-Status AvroScanner::_construct_row_without_jsonpath(const avro_value_t& avro_value, Chunk* chunk) {
+Status AvroScanner::_construct_row_without_jsonpath(const avro_value_t& avro_value, Chunk* chunk,
+                                                    const StreamMessageMeta* meta) {
     _found_columns.assign(chunk->num_columns(), false);
     size_t element_count = _data_idx_to_fieldname.size();
     avro_value_t element_value;
@@ -362,6 +391,12 @@ Status AvroScanner::_construct_row_without_jsonpath(const avro_value_t& avro_val
                 continue;
             }
             auto slot_desc = itr->second;
+            // A metadata alias is filled from the message meta in the null-fill pass, never from the
+            // payload; if a payload field uses the same name, skip it (cache as id = -1) so metadata wins.
+            if (_meta_col_by_slot_id.find(slot_desc->id()) != _meta_col_by_slot_id.end()) {
+                slot_info.id = -1;
+                continue;
+            }
             slot_info.id = slot_desc->id();
             // Store the intermediate avro type (not the destination type) so _construct_column
             // dispatches on the type the source column was built with. See _construct_row.
@@ -386,7 +421,13 @@ Status AvroScanner::_construct_row_without_jsonpath(const avro_value_t& avro_val
     for (int i = 0; i < _found_columns.size(); i++) {
         if (UNLIKELY(!_found_columns[i])) {
             auto* column = chunk->get_column_raw_ptr_by_index(i);
-            column->append_nulls(1);
+            // A hidden metadata slot always lands here (the by-name resolve above skips any same-named
+            // payload field); fill it from the message meta, NULL when the buffer carries none.
+            if (auto it = _meta_col_by_index.find(i); it != _meta_col_by_index.end()) {
+                RETURN_IF_ERROR(fill_stream_source_meta_column(it->second.kind, meta, column));
+            } else {
+                column->append_nulls(1);
+            }
         }
     }
     return Status::OK();

be/src/exec/file_scanner/avro_scanner.h

@@ -21,6 +21,7 @@
 #include "common/status.h"
 #include "exec/file_scanner/file_scanner.h"
 #include "exec/file_scanner/json_scanner.h"
+#include "exec/file_scanner/stream_source_meta.h"
 #include "exprs/json_functions.h"
 #include "fs/fs.h"
 #include "runtime/stream_load/load_stream_mgr.h"
@@ -39,6 +40,8 @@ namespace starrocks {
 
 using AvroPath = SimpleJsonPath;
 
+class StreamMessageMeta;
+
 class AvroScanner final : public FileScanner {
 public:
     AvroScanner(RuntimeState* state, RuntimeProfile* profile, const TBrokerScanRange& scan_range,
@@ -73,7 +76,7 @@ class AvroScanner final : public FileScanner {
     Status _create_src_chunk(ChunkPtr* chunk);
     Status _parse_avro(Chunk* chunk, const std::shared_ptr<SequentialFile>& file);
     void _report_error(const std::string& line, const std::string& err_msg);
-    Status _construct_row(const avro_value_t& avro_value, Chunk* chunk);
+    Status _construct_row(const avro_value_t& avro_value, Chunk* chunk, const StreamMessageMeta* meta);
     void _materialize_src_chunk_adaptive_nullable_column(ChunkPtr& chunk);
     Status _construct_column(const avro_value_t& input_value, Column* column, const TypeDescriptor& type_desc,
                              const std::string& col_name);
@@ -82,7 +85,7 @@ class AvroScanner final : public FileScanner {
     Status _handle_union(const avro_value_t* input_value, avro_value_t* branch);
     Status _get_array_element(const avro_value_t* cur_value, size_t idx, avro_value_t* element);
     std::string _preprocess_jsonpaths(std::string jsonpath);
-    Status _construct_row_without_jsonpath(const avro_value_t& avro_value, Chunk* chunk);
+    Status _construct_row_without_jsonpath(const avro_value_t& avro_value, Chunk* chunk, const StreamMessageMeta* meta);
 
     const TBrokerScanRange& _scan_range;
     serdes_t* _serdes;
@@ -99,13 +102,26 @@ class AvroScanner final : public FileScanner {
     std::unordered_map<std::string_view, SlotDescriptor*> _slot_desc_dict;
     // Maps each source slot id to its intermediate avro load type (see AvroScanner::_construct_avro_types).
     std::unordered_map<SlotId, TypeDescriptor> _slot_id_to_avro_type;
+    // Hidden source-metadata slots (routine load), filled from the message's ByteBuffer meta rather than
+    // the avro payload. _meta_col_by_slot_id keys by source slot id (the jsonpath path iterates slots);
+    // _meta_col_by_index keys by chunk column index for the by-name null-fill path. Both empty otherwise.
+    StreamSourceMetaColumns _meta_col_by_slot_id;
+    std::unordered_map<int, TRoutineLoadMetaColumn> _meta_col_by_index;
     std::vector<bool> _found_columns;
     std::vector<SlotInfo> _data_idx_to_slot;
     std::vector<std::string> _data_idx_to_fieldname;
     bool _init_data_idx_to_slot_once;
 
 #if BE_TEST
+public:
+    // Test-only: inject the per-message source metadata that production reads from the pipe buffer.
+    // BE_TEST reads avro from a file rather than the Kafka/Pulsar pipe, so the metadata-column fill path
+    // is otherwise unreachable from a test; this lets AvroScannerTest exercise it.
+    void set_test_stream_meta(const StreamMessageMeta* meta) { _test_meta = meta; }
+
+private:
     avro_file_reader_t _dbreader;
+    const StreamMessageMeta* _test_meta = nullptr;
 #endif
 };
 

be/src/exec/file_scanner/json_scanner.cpp

@@ -24,6 +24,7 @@
 #include "column/chunk.h"
 #include "column/column_helper.h"
 #include "column/vectorized_fwd.h"
+#include "exec/file_scanner/stream_source_meta.h"
 #include "exec/json_parser.h"
 #include "exprs/cast_expr.h"
 #include "exprs/column_ref.h"
@@ -287,6 +288,7 @@ JsonReader::JsonReader(RuntimeState* state, ScannerCounter* counter, JsonScanner
           _type_descs(std::move(type_descs)),
           _op_col_index(-1),
           _range_desc(range_desc) {
+    _meta_col_by_slot_id = build_stream_source_meta_columns(_scanner->_scan_range.params.stream_source_meta_columns);
     int index = 0;
     for (size_t i = 0; i < _slot_descs.size(); ++i) {
         const auto& desc = _slot_descs[i];
@@ -295,6 +297,8 @@ JsonReader::JsonReader(RuntimeState* state, ScannerCounter* counter, JsonScanner
         }
         if (UNLIKELY(desc->col_name() == "__op")) {
             _op_col_index = index;
+        } else if (auto m = _meta_col_by_slot_id.find(desc->id()); m != _meta_col_by_slot_id.end()) {
+            _meta_col_by_index.emplace(index, m->second);
         }
         index++;
         _slot_desc_dict.emplace(desc->col_name(), desc);
@@ -479,6 +483,15 @@ Status JsonReader::_read_rows(Chunk* chunk, int32_t rows_to_read, int32_t* rows_
 Status JsonReader::_construct_row_without_jsonpath(simdjson::ondemand::object* row, Chunk* chunk) {
     _parsed_columns.assign(chunk->num_columns(), false);
 
+    // Hidden source-metadata columns are filled from the buffer's Kafka/Pulsar metadata (routine load),
+    // never from the payload. If a payload field uses the same name as a metadata alias, the by-name
+    // resolve below skips it so the metadata value wins.
+#if BE_TEST
+    const StreamMessageMeta* meta = _scanner->_test_meta;
+#else
+    const StreamMessageMeta* meta = stream_source_meta_of(_file_stream_buffer);
+#endif
+
     faststring buffer;
     try {
         uint32_t key_index = 0;
@@ -530,8 +543,22 @@ Status JsonReader::_construct_row_without_jsonpath(simdjson::ondemand::object* r
                 auto slot_desc = itr->second;
                 auto type_desc = _type_desc_dict[key];
 
-                // update the prev parsed position
                 column_index = chunk->get_index_by_slot_id(slot_desc->id());
+                // A hidden metadata slot is filled from the message meta in the null-fill pass, never
+                // from the payload; if a payload field uses the same name as the metadata alias, skip it
+                // (cache as column_index = -1) so the metadata value still wins.
+                if (_meta_col_by_index.count(column_index) > 0) {
+                    if (_prev_parsed_position.size() <= key_index) {
+                        _prev_parsed_position.emplace_back(key);
+                    } else {
+                        _prev_parsed_position[key_index].key = key;
+                        _prev_parsed_position[key_index].column_index = -1;
+                    }
+                    key_index++;
+                    continue;
+                }
+
+                // update the prev parsed position
                 if (_prev_parsed_position.size() <= key_index) {
                     _prev_parsed_position.emplace_back(key, column_index, type_desc);
                 } else {
@@ -566,7 +593,7 @@ Status JsonReader::_construct_row_without_jsonpath(simdjson::ondemand::object* r
         return Status::DataQualityError(err_msg);
     }
 
-    // append null to the column without data.
+    // append null (or synthetic metadata) to the columns without data.
     for (int i = 0; i < chunk->num_columns(); i++) {
         if (!_parsed_columns[i]) {
             auto* column = chunk->get_column_raw_ptr_by_index(i);
@@ -577,6 +604,9 @@ Status JsonReader::_construct_row_without_jsonpath(simdjson::ondemand::object* r
                 } else {
                     column->append_datum(Datum((uint8_t)0));
                 }
+            } else if (auto it = _meta_col_by_index.find(i); it != _meta_col_by_index.end()) {
+                // Fill from meta (NULL when the buffer carries none).
+                RETURN_IF_ERROR(fill_stream_source_meta_column(it->second.kind, meta, column));
             } else {
                 column->append_nulls(1);
             }
@@ -588,6 +618,15 @@ Status JsonReader::_construct_row_without_jsonpath(simdjson::ondemand::object* r
 Status JsonReader::_construct_row_with_jsonpath(simdjson::ondemand::object* row, Chunk* chunk) {
     size_t slot_size = _slot_descs.size();
     size_t jsonpath_size = _scanner->_json_paths.size();
+#if BE_TEST
+    const StreamMessageMeta* meta = _scanner->_test_meta;
+#else
+    const StreamMessageMeta* meta = stream_source_meta_of(_file_stream_buffer);
+#endif
+    // Hidden metadata columns carry no jsonpath, so they are transparent to the positional
+    // slot->jsonpath mapping: a jsonpath maps to the i-th non-metadata slot. This keeps payload columns
+    // aligned even when a metadata column sits before jsonpath-backed columns.
+    size_t meta_count = 0;
     for (size_t i = 0; i < slot_size; i++) {
         if (_slot_descs[i] == nullptr) {
             continue;
@@ -596,7 +635,18 @@ Status JsonReader::_construct_row_with_jsonpath(simdjson::ondemand::object* row,
 
         // The columns in JsonReader's chunk are all in NullableColumn type;
         auto* column = down_cast<NullableColumn*>(chunk->get_column_raw_ptr_by_slot_id(_slot_descs[i]->id()));
-        if (i >= jsonpath_size) {
+
+        // Hidden source-metadata slots are filled from the message meta, not the payload, before the
+        // jsonpath extraction. Fill unconditionally (NULL when the buffer carries no meta) and always
+        // advance meta_count, so a metadata slot stays transparent to the positional jsonpath mapping
+        // regardless of whether meta is present.
+        if (auto it = _meta_col_by_slot_id.find(_slot_descs[i]->id()); UNLIKELY(it != _meta_col_by_slot_id.end())) {
+            RETURN_IF_ERROR(fill_stream_source_meta_column(it->second.kind, meta, column));
+            meta_count++;
+            continue;
+        }
+        size_t jp = i - meta_count;
+        if (jp >= jsonpath_size) {
             if (strcmp(column_name, "__op") == 0) {
                 // special treatment for __op column, fill default value '0' rather than null
                 if (column->is_binary()) {
@@ -616,7 +666,7 @@ Status JsonReader::_construct_row_with_jsonpath(simdjson::ondemand::object* row,
         // simdjson's api is limited, which coult not convert ondemand::object to ondemand::value.
         // As a workaround, extract procedure is duplicated, for both ondemand::object and ondemand::value
         // TODO(mofei) make it more elegant
-        if (_scanner->_json_paths[i].size() == 1 && _scanner->_json_paths[i][0].key == "$") {
+        if (_scanner->_json_paths[jp].size() == 1 && _scanner->_json_paths[jp][0].key == "$") {
             // add_nullable_column may invoke a for-range iterating to the row.
             // If the for-range iterating is invoked after field access, or a second for-range iterating is invoked,
             // it would get an error "Objects and arrays can only be iterated when they are first encountered",
@@ -626,7 +676,7 @@ Status JsonReader::_construct_row_with_jsonpath(simdjson::ondemand::object* row,
                     column, _slot_descs[i]->type(), _slot_descs[i]->col_name(), row, !_strict_mode));
         } else {
             simdjson::ondemand::value val;
-            auto st = JsonFunctions::extract_from_object(*row, _scanner->_json_paths[i], &val);
+            auto st = JsonFunctions::extract_from_object(*row, _scanner->_json_paths[jp], &val);
             if (st.ok()) {
                 RETURN_IF_ERROR(_construct_column(val, column, _slot_descs[i]->type(), _slot_descs[i]->col_name()));
             } else if (st.is_not_found()) {