wangyong9999 commented on code in PR #245: URL: https://github.com/apache/paimon-cpp/pull/245#discussion_r3850610606
########## src/paimon/core/index/pksorted/pk_sorted_data_file_reader.cpp: ########## @@ -0,0 +1,173 @@ +/* + * 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 "paimon/core/index/pksorted/pk_sorted_data_file_reader.h" + +#include <limits> +#include <map> +#include <optional> +#include <string> +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "fmt/format.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/operation/internal_read_context.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/read_context.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/reader/file_batch_reader.h" + +namespace paimon { + +Result<std::unique_ptr<PkSortedDataFileReader>> PkSortedDataFileReader::Create( + const std::string& root_path, const std::shared_ptr<TableSchema>& table_schema, + int32_t field_id, const std::shared_ptr<FileStorePathFactory>& path_factory, + const std::string& branch, const CoreOptions& options, + const std::shared_ptr<Executor>& executor, const std::shared_ptr<MemoryPool>& pool) { + std::map<std::string, std::string> read_options = options.ToMap(); + read_options[Options::BRANCH] = branch; + ReadContextBuilder builder(root_path); + builder.SetReadFieldIds({field_id}) + .SetOptions(read_options) + .WithBranch(branch) + .WithFileSystem(options.GetFileSystem()) + .WithExecutor(executor) + .WithMemoryPool(pool) + .EnablePrefetch(false) + .EnablePredicateFilter(false); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context, builder.Finish()); + auto shared_read_context = std::shared_ptr<ReadContext>(std::move(read_context)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<InternalReadContext> internal_context, + InternalReadContext::Create(shared_read_context, table_schema, read_options)); + auto shared_internal_context = + std::shared_ptr<InternalReadContext>(std::move(internal_context)); + return std::unique_ptr<PkSortedDataFileReader>( + new PkSortedDataFileReader(path_factory, shared_internal_context, pool, executor)); +} + +PkSortedDataFileReader::PkSortedDataFileReader( + const std::shared_ptr<FileStorePathFactory>& path_factory, + const std::shared_ptr<InternalReadContext>& context, const std::shared_ptr<MemoryPool>& pool, + const std::shared_ptr<Executor>& executor) + : RawFileSplitRead(path_factory, context, pool, executor) {} + +Status PkSortedDataFileReader::ReadFile(const BinaryRow& partition, int32_t bucket, + const std::shared_ptr<DataFileMeta>& file, + const BatchConsumer& consumer) const { + if (file == nullptr) { + return Status::Invalid("Primary-key sorted-index source file is null."); + } + if (file->row_count < 0) { + return Status::Invalid(fmt::format("Source file {} has negative row count {}.", + file->file_name, file->row_count)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<DataFilePathFactory> data_file_path_factory, + path_factory_->CreateDataFilePathFactory(partition, bucket)); + PAIMON_ASSIGN_OR_RAISE( + std::vector<std::unique_ptr<FileBatchReader>> readers, + CreateRawFileReaders(partition, {file}, raw_read_schema_, /*predicate=*/nullptr, + /*dv_factory=*/{}, /*row_ranges=*/std::nullopt, data_file_path_factory, + /*extra_format_options=*/{})); + if (readers.size() != 1) { + return Status::Invalid( + fmt::format("Expected one physical reader for source file {}, but got {}.", + file->file_name, readers.size())); + } + std::unique_ptr<FileBatchReader> reader = std::move(readers[0]); + ScopeGuard close_guard([&]() { reader->Close(); }); + PAIMON_ASSIGN_OR_RAISE(uint64_t physical_row_count, reader->GetNumberOfRows()); + if (physical_row_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) || + static_cast<int64_t>(physical_row_count) != file->row_count) { + return Status::Invalid(fmt::format( + "Physical row count {} of source file {} does not match metadata row count {}.", + physical_row_count, file->file_name, file->row_count)); + } + + int64_t rows_read = 0; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + break; + } + auto& [batch, bitmap] = batch_with_bitmap; + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (array == nullptr || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid( + fmt::format("Source file {} did not return a struct batch.", file->file_name)); + } + auto struct_array = checked_pointer_cast<arrow::StructArray>(array); + if (struct_array->num_fields() != 1) { + return Status::Invalid( + fmt::format("Source file {} returned {} fields for a single-column index build.", + file->file_name, struct_array->num_fields())); + } + if (static_cast<int64_t>(bitmap.Cardinality()) != struct_array->length()) { + return Status::Invalid( + fmt::format("Source file {} was filtered while building a physical-row index.", + file->file_name)); + } + std::vector<int64_t> positions; + positions.reserve(static_cast<size_t>(struct_array->length())); + for (int64_t index = 0; index < struct_array->length(); ++index) { + PAIMON_ASSIGN_OR_RAISE(uint64_t physical_position, + reader->GetPreviousBatchFileRowId(static_cast<uint64_t>(index))); + if (physical_position > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) || + static_cast<int64_t>(physical_position) != rows_read + index) { + return Status::Invalid(fmt::format( + "Source file {} returned non-contiguous physical row position {} at row {}.", + file->file_name, physical_position, rows_read + index)); + } + positions.push_back(static_cast<int64_t>(physical_position)); + } + PAIMON_RETURN_NOT_OK(consumer(struct_array, positions)); Review Comment: `positions` never reaches a consumer that uses it. The loop above already asserts `physical_position == rows_read + index` for every row, and `physical_row_count == file->row_count`, `bitmap.Cardinality() == length`, and the final `rows_read == file->row_count` pin the same property from three other directions — so the vector is always exactly `[rows_read, rows_read + length)`. The one consumer (`PkSortedIndexBuilder::Build`) only compares its `size()` and derives the ordinal from its own `rows_buffered + index`. So this materializes an int64 vector per batch and widens `BatchConsumer` to two parameters to carry a value nobody reads. Either use it for the ordinal (`file_base + positions[i]`, which is what the group-row-id contract suggests) or drop the parameter and keep the contiguity assertion local to the reader. ########## src/paimon/core/index/pksorted/pk_sorted_index_builder.cpp: ########## @@ -0,0 +1,276 @@ +/* + * 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 "paimon/core/index/pksorted/pk_sorted_index_builder.h" + +#include <algorithm> +#include <limits> +#include <map> +#include <string> +#include <utility> + +#include "arrow/api.h" +#include "arrow/array/concatenate.h" +#include "arrow/c/bridge.h" +#include "fmt/format.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/casting/casting_utils.h" +#include "paimon/core/global_index/global_index_file_manager.h" +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/core/index/pk/primary_key_index_source_policy.h" +#include "paimon/core/index/pksorted/pk_sorted_data_file_reader.h" +#include "paimon/core/index/pksorted/pk_sorted_index_file.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h" +#include "paimon/core/mergetree/external_sort_buffer.h" +#include "paimon/core/mergetree/in_memory_sort_buffer.h" +#include "paimon/core/mergetree/sort_buffer.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/fs/file_system.h" +#include "paimon/global_index/io/global_index_file_writer.h" +#include "paimon/record_batch.h" + +namespace paimon { +namespace { + +constexpr char kRowIdFieldName[] = "_PK_INDEX_ROW_ID"; + +class TrackingGlobalIndexFileWriter : public GlobalIndexFileWriter { + public: + explicit TrackingGlobalIndexFileWriter(const std::shared_ptr<GlobalIndexFileManager>& delegate) + : delegate_(delegate) {} + + Result<std::string> NewFileName(const std::string& prefix) const override { + PAIMON_ASSIGN_OR_RAISE(std::string file_name, delegate_->NewFileName(prefix)); + created_file_names_.push_back(file_name); + return file_name; + } + + Result<std::unique_ptr<OutputStream>> NewOutputStream( + const std::string& file_name) const override { + return delegate_->NewOutputStream(file_name); + } + + Result<int64_t> GetFileSize(const std::string& file_name) const override { + return delegate_->GetFileSize(file_name); + } + + std::string ToPath(const std::string& file_name) const override { + return delegate_->ToPath(file_name); + } + + void Cleanup(const std::shared_ptr<FileSystem>& fs) const { + for (const std::string& file_name : created_file_names_) { + [[maybe_unused]] Status status = fs->Delete(delegate_->ToPath(file_name)); + } + } + + private: + std::shared_ptr<GlobalIndexFileManager> delegate_; + mutable std::vector<std::string> created_file_names_; +}; + +} // namespace + +Result<std::unique_ptr<PkSortedIndexBuilder>> PkSortedIndexBuilder::Create( + const std::string& root_path, const std::string& branch, const BinaryRow& partition, + int32_t bucket, const std::shared_ptr<TableSchema>& table_schema, + const PrimaryKeyIndexDefinition& definition, + const std::shared_ptr<FileStorePathFactory>& path_factory, const CoreOptions& options, + const std::shared_ptr<IOManager>& io_manager, bool enable_multi_thread_spill, + const std::shared_ptr<Executor>& executor, const std::shared_ptr<MemoryPool>& pool) { + if (definition.GetFamily() != PrimaryKeyIndexDefinition::Family::BTREE) { + return Status::Invalid("PkSortedIndexBuilder only supports BTree definitions."); + } + PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(definition.FieldId())); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<PkSortedDataFileReader> data_file_reader, + PkSortedDataFileReader::Create(root_path, table_schema, definition.FieldId(), path_factory, + branch, options, executor, pool)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<IndexPathFactory> index_path_factory, + path_factory->CreateIndexFileFactory(partition, bucket)); + return std::unique_ptr<PkSortedIndexBuilder>(new PkSortedIndexBuilder( + partition, bucket, std::move(field), definition, + std::shared_ptr<PkSortedDataFileReader>(std::move(data_file_reader)), + options.GetFileSystem(), std::shared_ptr<IndexPathFactory>(std::move(index_path_factory)), + options, io_manager, enable_multi_thread_spill, pool)); +} + +Result<std::shared_ptr<IndexFileMeta>> PkSortedIndexBuilder::Build( + const std::vector<std::shared_ptr<DataFileMeta>>& source_files) const { + if (source_files.empty()) { + return Status::Invalid("Cannot build a sorted index for an empty data level."); + } + for (const std::shared_ptr<DataFileMeta>& file : source_files) { + if (file == nullptr) { + return Status::Invalid("A sorted index source file is null."); + } + } + std::vector<std::shared_ptr<DataFileMeta>> ordered_files = source_files; + std::sort( + ordered_files.begin(), ordered_files.end(), + [](const std::shared_ptr<DataFileMeta>& left, const std::shared_ptr<DataFileMeta>& right) { + return left->file_name < right->file_name; + }); + int32_t data_level = ordered_files.front()->level; + std::vector<PrimaryKeyIndexSourceFile> source_metas; + source_metas.reserve(ordered_files.size()); + for (const std::shared_ptr<DataFileMeta>& file : ordered_files) { + if (file == nullptr || file->level != data_level || + !PrimaryKeyIndexSourcePolicy::ShouldRead(*file)) { + return Status::Invalid( + "A sorted index can only cover compacted files from one positive data level."); + } + source_metas.emplace_back(file->file_name, file->row_count); + } + + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FieldsComparator> unique_comparator, + FieldsComparator::Create({field_}, /*is_ascending_order=*/true)); + auto comparator = std::shared_ptr<FieldsComparator>(std::move(unique_comparator)); + DataField row_id_field(std::numeric_limits<int32_t>::max(), Review Comment: This column carries the same number the sort buffer already assigns. With `last_sequence_number = -1` and rows written in group order, `KeyValue.sequence_number` *is* the group row id — that is exactly what `BuildFromSortedReader` reads back from `record_batch->column(0)`. The cost of the duplicate: an extra int64 per row, which doubles the sort-buffer footprint for a 4-byte indexed type on the one path that hard-fails when the buffer quota is exhausted, plus three comparators (`comparator`, `sequence_comparator`, `in_memory_comparator`) where one would do. Dropping it should work directly: `stable_sort` on the value alone keeps insertion order for ties, and `SortMergeReaderWithMinHeap` already falls back to `sequence_number` when the user-defined sequence comparator is null. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
