feat(lumina): support Lumina tag filtering

Author: lxy-9602

include/paimon/global_index/global_index_io_meta.h

@@ -16,8 +16,11 @@
 
 #pragma once
 
+#include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
+#include <vector>
 
 #include "paimon/memory/bytes.h"
 #include "paimon/utils/range.h"
@@ -27,14 +30,24 @@ namespace paimon {
 struct PAIMON_EXPORT GlobalIndexIOMeta {
     GlobalIndexIOMeta(const std::string& _file_path, int64_t _file_size,
                       const std::shared_ptr<Bytes>& _metadata)
-        : file_path(_file_path), file_size(_file_size), metadata(_metadata) {}
+        : GlobalIndexIOMeta(_file_path, _file_size, _metadata, std::nullopt) {}
+
+    GlobalIndexIOMeta(const std::string& _file_path, int64_t _file_size,
+                      const std::shared_ptr<Bytes>& _metadata,
+                      const std::optional<std::vector<int32_t>>& _extra_field_ids)
+        : file_path(_file_path),
+          file_size(_file_size),
+          metadata(_metadata),
+          extra_field_ids(_extra_field_ids) {}
 
     std::string file_path;
     int64_t file_size;
     /// Optional binary metadata associated with the file, such as serialized
     /// secondary index structures or inline index bytes.
     /// May be null if no additional metadata is available.
     std::shared_ptr<Bytes> metadata;
+    /// Optional table field ids materialized together with the indexed field.
+    std::optional<std::vector<int32_t>> extra_field_ids;
 };
 
 }  // namespace paimon

include/paimon/global_index/global_indexer.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -36,6 +37,11 @@ class PAIMON_EXPORT GlobalIndexer {
  public:
     virtual ~GlobalIndexer() = default;
 
+    /// Returns additional table fields required during index construction.
+    virtual Result<std::optional<std::vector<std::string>>> GetExtraFieldNames() const {
+        return std::optional<std::vector<std::string>>(std::nullopt);
+    }
+
     /// Creates a writer for building a global index on a specific field.
     ///
     /// @param field_name     Name of the field to be indexed.

src/paimon/core/global_index/global_index_scan_impl.cpp

@@ -20,6 +20,7 @@
 #include <utility>
 
 #include "arrow/c/bridge.h"
+#include "arrow/c/helpers.h"
 #include "paimon/common/global_index/offset_global_index_reader.h"
 #include "paimon/common/global_index/union_global_index_reader.h"
 #include "paimon/common/utils/scope_guard.h"
@@ -185,7 +186,6 @@ Result<std::vector<std::shared_ptr<GlobalIndexReader>>> GlobalIndexScanImpl::Cre
             if (row_range_index && !row_range_index->Intersects(range.from, range.to)) {
                 continue;
             }
-            // TODO(xinyu.lxy): c_arrow_schema may contains additional associated fields.
             auto arrow_field = DataField::ConvertDataFieldToArrowField(field);
             auto arrow_schema = arrow::schema({arrow_field});
 
@@ -223,7 +223,7 @@ GlobalIndexIOMeta GlobalIndexScanImpl::ToGlobalIndexIOMeta(
     assert(index_meta->GetGlobalIndexMeta());
     const auto& global_index_meta = index_meta->GetGlobalIndexMeta().value();
     return {index_file_manager_->ToPath(index_meta), index_meta->FileSize(),
-            global_index_meta.index_meta};
+            global_index_meta.index_meta, global_index_meta.extra_field_ids};
 }
 
 Result<std::shared_ptr<GlobalIndexResult>> GlobalIndexScanImpl::Scan(

src/paimon/core/global_index/global_index_write_task.cpp

@@ -16,7 +16,10 @@
 
 #include "paimon/global_index/global_index_write_task.h"
 
+#include <set>
+
 #include "arrow/c/bridge.h"
+#include "arrow/c/helpers.h"
 #include "paimon/common/table/special_fields.h"
 #include "paimon/common/types/data_field.h"
 #include "paimon/common/utils/arrow/status_utils.h"
@@ -35,6 +38,17 @@
 #include "paimon/table/source/table_read.h"
 namespace paimon {
 namespace {
+Result<std::unique_ptr<GlobalIndexer>> CreateGlobalIndexer(const std::string& index_type,
+                                                           const CoreOptions& core_options) {
+    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<GlobalIndexer> indexer,
+                           GlobalIndexerFactory::Get(index_type, core_options.ToMap()));
+    if (!indexer) {
+        return Status::Invalid(
+            fmt::format("Unknown index type {}, may not registered", index_type));
+    }
+    return indexer;
+}
+
 Result<std::shared_ptr<GlobalIndexFileManager>> CreateGlobalIndexFileManager(
     const std::string& table_path, const std::shared_ptr<TableSchema>& table_schema,
     const CoreOptions& core_options, const std::shared_ptr<MemoryPool>& pool) {
@@ -58,43 +72,95 @@ Result<std::shared_ptr<GlobalIndexFileManager>> CreateGlobalIndexFileManager(
 }
 
 Result<std::shared_ptr<GlobalIndexWriter>> CreateGlobalIndexWriter(
-    const std::string& index_type, const DataField& field,
+    const GlobalIndexer& indexer, const DataField& field,
+    const std::vector<DataField>& extra_fields,
     const std::shared_ptr<GlobalIndexFileManager>& index_file_manager,
-    const CoreOptions& core_options, const std::shared_ptr<MemoryPool>& pool) {
-    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<GlobalIndexer> indexer,
-                           GlobalIndexerFactory::Get(index_type, core_options.ToMap()));
-    if (!indexer) {
-        return Status::Invalid(
-            fmt::format("Unknown index type {}, may not registered", index_type));
+    const std::shared_ptr<MemoryPool>& pool) {
+    arrow::FieldVector arrow_fields;
+    arrow_fields.reserve(extra_fields.size() + 1);
+    arrow_fields.push_back(DataField::ConvertDataFieldToArrowField(field));
+    for (const auto& extra_field : extra_fields) {
+        arrow_fields.push_back(DataField::ConvertDataFieldToArrowField(extra_field));
     }
-    // TODO(xinyu.lxy): may add additional fields to read for index write
-    auto arrow_field = DataField::ConvertDataFieldToArrowField(field);
-    auto arrow_schema = arrow::schema({arrow_field});
+    auto arrow_schema = arrow::schema(arrow_fields);
     ArrowSchema c_arrow_schema;
     PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow_schema, &c_arrow_schema));
-    return indexer->CreateWriter(field.Name(), &c_arrow_schema, index_file_manager, pool);
+    ScopeGuard guard([&]() { ArrowSchemaRelease(&c_arrow_schema); });
+    return indexer.CreateWriter(field.Name(), &c_arrow_schema, index_file_manager, pool);
+}
+
+Result<std::vector<DataField>> GetExtraFields(const TableSchema& table_schema,
+                                              const std::string& field_name,
+                                              const std::vector<std::string>& extra_field_names) {
+    std::vector<DataField> extra_fields;
+    extra_fields.reserve(extra_field_names.size());
+    std::set<std::string> dedup_field_names;
+    for (const auto& extra_field_name : extra_field_names) {
+        if (extra_field_name == field_name || !dedup_field_names.insert(extra_field_name).second) {
+            continue;
+        }
+        PAIMON_ASSIGN_OR_RAISE(DataField extra_field, table_schema.GetField(extra_field_name));
+        extra_fields.push_back(extra_field);
+    }
+    return extra_fields;
+}
+
+std::vector<std::string> BuildReadFieldNames(const std::string& field_name,
+                                             const std::vector<DataField>& extra_fields) {
+    std::vector<std::string> read_field_names;
+    read_field_names.reserve(extra_fields.size() + 2);
+    read_field_names.push_back(field_name);
+    for (const auto& extra_field : extra_fields) {
+        read_field_names.push_back(extra_field.Name());
+    }
+    read_field_names.push_back(SpecialFields::RowId().Name());
+    return read_field_names;
+}
+
+std::vector<std::string> BuildWriterFieldNames(const std::string& field_name,
+                                               const std::vector<DataField>& extra_fields) {
+    std::vector<std::string> writer_field_names;
+    writer_field_names.reserve(extra_fields.size() + 1);
+    writer_field_names.push_back(field_name);
+    for (const auto& extra_field : extra_fields) {
+        writer_field_names.push_back(extra_field.Name());
+    }
+    return writer_field_names;
+}
+
+std::optional<std::vector<int32_t>> GetExtraFieldIds(const std::vector<DataField>& extra_fields) {
+    if (extra_fields.empty()) {
+        return std::nullopt;
+    }
+    std::vector<int32_t> extra_field_ids;
+    extra_field_ids.reserve(extra_fields.size());
+    for (const auto& extra_field : extra_fields) {
+        extra_field_ids.push_back(extra_field.Id());
+    }
+    return extra_field_ids;
 }
 
 Result<std::unique_ptr<BatchReader>> CreateBatchReader(
-    const std::string& table_path, const std::string& field_name,
+    const std::string& table_path, const std::vector<std::string>& read_field_names,
     const std::shared_ptr<IndexedSplit>& indexed_split, const CoreOptions& core_options,
     const std::shared_ptr<MemoryPool>& pool) {
     ReadContextBuilder read_context_builder(table_path);
     read_context_builder.SetOptions(core_options.ToMap())
         .WithFileSystem(core_options.GetFileSystem())
         .EnablePrefetch(true)
         .WithMemoryPool(pool)
-        .SetReadFieldNames({field_name, SpecialFields::RowId().Name()});
+        .SetReadFieldNames(read_field_names);
     PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context,
                            read_context_builder.Finish());
     PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> table_read,
                            TableRead::Create(std::move(read_context)));
     return table_read->CreateReader(indexed_split);
 }
 
-Result<std::vector<GlobalIndexIOMeta>> BuildIndex(const std::string& field_name, const Range& range,
-                                                  BatchReader* batch_reader,
-                                                  GlobalIndexWriter* global_index_writer) {
+Result<std::vector<GlobalIndexIOMeta>> BuildIndex(
+    const std::string& field_name, const Range& range,
+    const std::vector<std::string>& writer_field_names, BatchReader* batch_reader,
+    GlobalIndexWriter* global_index_writer) {
     while (true) {
         PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch read_batch, batch_reader->NextBatch());
         if (BatchReader::IsEofBatch(read_batch)) {
@@ -131,8 +197,20 @@ Result<std::vector<GlobalIndexIOMeta>> BuildIndex(const std::string& field_name,
             }
             relative_row_ids.push_back(row_id - range.from);
         }
-        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray> new_array,
-                                          arrow::StructArray::Make({indexed_array}, {field_name}));
+        std::vector<std::shared_ptr<arrow::Array>> writer_arrays;
+        writer_arrays.reserve(writer_field_names.size());
+        for (const auto& writer_field_name : writer_field_names) {
+            auto writer_array = struct_array->GetFieldByName(writer_field_name);
+            if (!writer_array) {
+                return Status::Invalid(
+                    fmt::format("read array does not contain {} field in GlobalIndexWriteTask",
+                                writer_field_name));
+            }
+            writer_arrays.push_back(writer_array);
+        }
+        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+            std::shared_ptr<arrow::StructArray> new_array,
+            arrow::StructArray::Make(writer_arrays, writer_field_names));
         ::ArrowArray c_new_array;
         PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*new_array, &c_new_array));
         PAIMON_RETURN_NOT_OK(
@@ -157,8 +235,8 @@ Result<std::shared_ptr<CommitMessage>> ToCommitMessage(
         index_file_metas.push_back(std::make_shared<IndexFileMeta>(
             index_type, PathUtil::GetName(io_meta.file_path), io_meta.file_size, range.Count(),
             /*dv_ranges=*/std::nullopt, external_path,
-            GlobalIndexMeta(range.from, range.to, field_id,
-                            /*extra_field_ids=*/std::nullopt, io_meta.metadata)));
+            GlobalIndexMeta(range.from, range.to, field_id, io_meta.extra_field_ids,
+                            io_meta.metadata)));
     }
     DataIncrement data_increment(std::move(index_file_metas));
     return std::make_shared<CommitMessageImpl>(partition, bucket,
@@ -199,6 +277,17 @@ Result<std::shared_ptr<CommitMessage>> GlobalIndexWriteTask::WriteIndex(
     }
     PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options,
                            CoreOptions::FromMap(final_options, file_system));
+    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<GlobalIndexer> indexer,
+                           CreateGlobalIndexer(index_type, core_options));
+    PAIMON_ASSIGN_OR_RAISE(std::optional<std::vector<std::string>> extra_field_names,
+                           indexer->GetExtraFieldNames());
+    PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(field_name));
+    PAIMON_ASSIGN_OR_RAISE(std::vector<DataField> extra_fields,
+                           GetExtraFields(*table_schema, field_name,
+                                          extra_field_names.value_or(std::vector<std::string>())));
+    std::optional<std::vector<int32_t>> extra_field_ids = GetExtraFieldIds(extra_fields);
+    std::vector<std::string> writer_field_names = BuildWriterFieldNames(field_name, extra_fields);
+    std::vector<std::string> read_field_names = BuildReadFieldNames(field_name, extra_fields);
 
     // create index file manager
     PAIMON_ASSIGN_OR_RAISE(
@@ -208,23 +297,25 @@ Result<std::shared_ptr<CommitMessage>> GlobalIndexWriteTask::WriteIndex(
     // create batch reader
     PAIMON_ASSIGN_OR_RAISE(
         std::unique_ptr<BatchReader> batch_reader,
-        CreateBatchReader(table_path, field_name, indexed_split, core_options, pool));
+        CreateBatchReader(table_path, read_field_names, indexed_split, core_options, pool));
 
     // create global index writer
-    PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(field_name));
     PAIMON_ASSIGN_OR_RAISE(
         std::shared_ptr<GlobalIndexWriter> global_index_writer,
-        CreateGlobalIndexWriter(index_type, field, index_file_manager, core_options, pool));
+        CreateGlobalIndexWriter(*indexer, field, extra_fields, index_file_manager, pool));
 
     ScopeGuard guard([&]() {
         global_index_writer.reset();
         batch_reader.reset();
     });
 
     // read from data split and write to index writer
-    PAIMON_ASSIGN_OR_RAISE(
-        std::vector<GlobalIndexIOMeta> global_index_io_metas,
-        BuildIndex(field_name, range, batch_reader.get(), global_index_writer.get()));
+    PAIMON_ASSIGN_OR_RAISE(std::vector<GlobalIndexIOMeta> global_index_io_metas,
+                           BuildIndex(field_name, range, writer_field_names, batch_reader.get(),
+                                      global_index_writer.get()));
+    for (auto& io_meta : global_index_io_metas) {
+        io_meta.extra_field_ids = extra_field_ids;
+    }
 
     // generate commit message
     return ToCommitMessage(index_type, field.Id(), range, global_index_io_metas,