From 45247b01c8411e14baa87936e2e7470819775011 Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Mon, 17 Aug 2026 15:16:26 +0800 Subject: [PATCH 1/4] feat(file-index): support writing file indexes --- include/paimon/defs.h | 4 + include/paimon/file_index/file_index_format.h | 21 +++ src/paimon/CMakeLists.txt | 5 + src/paimon/common/defs.cpp | 1 + .../common/file_index/file_index_format.cpp | 109 ++++++++++++ .../file_index/file_index_format_test.cpp | 21 +++ .../common/io/byte_array_output_stream.cpp | 79 +++++++++ .../common/io/byte_array_output_stream.h | 69 ++++++++ .../io/byte_array_output_stream_test.cpp | 83 +++++++++ .../io/memory_segment_output_stream.cpp | 18 +- .../io/memory_segment_output_stream_test.cpp | 12 ++ .../core/append/append_only_writer_test.cpp | 37 ++++ src/paimon/core/core_options.cpp | 8 + src/paimon/core/core_options.h | 1 + src/paimon/core/core_options_test.cpp | 3 + .../io/append_data_file_writer_factory.cpp | 6 + src/paimon/core/io/data_file_index_writer.cpp | 164 ++++++++++++++++++ src/paimon/core/io/data_file_index_writer.h | 93 ++++++++++ .../core/io/data_file_index_writer_test.cpp | 161 +++++++++++++++++ src/paimon/core/io/data_file_writer.cpp | 25 +-- src/paimon/core/io/data_file_writer.h | 17 +- src/paimon/core/io/data_file_writer_base.h | 134 ++++++++++++++ .../core/io/data_file_writer_factory.cpp | 15 ++ src/paimon/core/io/data_file_writer_factory.h | 6 + src/paimon/core/io/file_index_options.cpp | 102 +++++++++++ src/paimon/core/io/file_index_options.h | 61 +++++++ .../core/io/key_value_data_file_writer.cpp | 30 +--- .../core/io/key_value_data_file_writer.h | 19 +- .../io/key_value_data_file_writer_factory.cpp | 6 + ...edding_append_data_file_writer_factory.cpp | 6 + ...ing_key_value_data_file_writer_factory.cpp | 6 + src/paimon/core/io/single_file_writer.h | 25 ++- 32 files changed, 1263 insertions(+), 84 deletions(-) create mode 100644 src/paimon/common/io/byte_array_output_stream.cpp create mode 100644 src/paimon/common/io/byte_array_output_stream.h create mode 100644 src/paimon/common/io/byte_array_output_stream_test.cpp create mode 100644 src/paimon/core/io/data_file_index_writer.cpp create mode 100644 src/paimon/core/io/data_file_index_writer.h create mode 100644 src/paimon/core/io/data_file_index_writer_test.cpp create mode 100644 src/paimon/core/io/data_file_writer_base.h create mode 100644 src/paimon/core/io/file_index_options.cpp create mode 100644 src/paimon/core/io/file_index_options.h diff --git a/include/paimon/defs.h b/include/paimon/defs.h index d1ebf507a..e944587f2 100644 --- a/include/paimon/defs.h +++ b/include/paimon/defs.h @@ -405,6 +405,10 @@ struct PAIMON_EXPORT Options { /// "file-index.read.enabled" - Whether enabled read file index. Default value is "true". static const char FILE_INDEX_READ_ENABLED[]; + /// "file-index.in-manifest-threshold" - The threshold to store file index bytes in the + /// manifest. Default value is 500B. + static const char FILE_INDEX_IN_MANIFEST_THRESHOLD[]; + /// "data-file.external-paths" - The external paths where the data of this table will be /// written, multiple elements separated by commas. static const char DATA_FILE_EXTERNAL_PATHS[]; diff --git a/include/paimon/file_index/file_index_format.h b/include/paimon/file_index/file_index_format.h index b46dee8c4..38459845a 100644 --- a/include/paimon/file_index/file_index_format.h +++ b/include/paimon/file_index/file_index_format.h @@ -19,6 +19,7 @@ #pragma once #include +#include #include #include #include @@ -32,6 +33,8 @@ struct ArrowSchema; namespace paimon { class InputStream; class MemoryPool; +class Bytes; +class OutputStream; /// Defines the on-disk format and versioning for Paimon file-level indexes. /// File index file format. Put all column and offset in the header. @@ -88,6 +91,8 @@ class MemoryPool; class PAIMON_EXPORT FileIndexFormat { public: class Reader; + class Writer; + using ColumnIndexes = std::map>>; /// Creates a `Reader` to parse a index file (may contain multiple indexes) from the given input /// stream. /// @@ -98,12 +103,28 @@ class PAIMON_EXPORT FileIndexFormat { static Result> CreateReader( const std::shared_ptr& input_stream, const std::shared_ptr& pool); + /// Creates a `Writer` which serializes a complete V1 file index container. + static Result> CreateWriter( + const std::shared_ptr& output_stream, + const std::shared_ptr& pool); + public: static const int64_t MAGIC; static const int32_t EMPTY_INDEX_FLAG; static const int32_t V_1; }; +/// Writer for file index file. +class FileIndexFormat::Writer { + public: + virtual ~Writer() = default; + + /// Writes all column indexes. This is a terminal, one-shot operation. + virtual Status WriteColumnIndexes(const FileIndexFormat::ColumnIndexes& indexes) = 0; + + virtual Status Close() = 0; +}; + /// Reader for file index file. class FileIndexFormat::Reader { public: diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 9b0807b64..daa39fb53 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -85,6 +85,7 @@ set(PAIMON_COMMON_SRCS common/global_index/global_indexer_factory.cpp common/io/buffered_input_stream.cpp common/io/byte_array_input_stream.cpp + common/io/byte_array_output_stream.cpp common/io/data_input_stream.cpp common/io/data_output_stream.cpp common/io/memory_segment_output_stream.cpp @@ -270,6 +271,8 @@ set(PAIMON_CORE_SRCS core/io/data_file_meta.cpp core/io/data_file_meta_serializer.cpp core/io/data_file_path_factory.cpp + core/io/data_file_index_writer.cpp + core/io/file_index_options.cpp core/io/append_data_file_writer_factory.cpp core/io/blob_data_file_writer_factory.cpp core/io/data_file_writer_factory.cpp @@ -577,6 +580,7 @@ if(PAIMON_BUILD_TESTS) common/global_index/rangebitmap/range_bitmap_global_index_test.cpp common/global_index/wrap/file_index_reader_wrapper_test.cpp common/io/byte_array_input_stream_test.cpp + common/io/byte_array_output_stream_test.cpp common/io/data_input_output_stream_test.cpp common/io/buffered_input_stream_test.cpp common/io/memory_segment_output_stream_test.cpp @@ -752,6 +756,7 @@ if(PAIMON_BUILD_TESTS) core/io/complete_row_tracking_fields_reader_test.cpp core/io/vector_file_batch_reader_test.cpp core/io/data_file_meta_test.cpp + core/io/data_file_index_writer_test.cpp core/io/file_index_evaluator_test.cpp core/io/single_file_writer_test.cpp core/io/rolling_blob_file_writer_test.cpp diff --git a/src/paimon/common/defs.cpp b/src/paimon/common/defs.cpp index bac4f16f7..ef35940e9 100644 --- a/src/paimon/common/defs.cpp +++ b/src/paimon/common/defs.cpp @@ -100,6 +100,7 @@ const char Options::PARTIAL_UPDATE_REMOVE_RECORD_ON_SEQUENCE_GROUP[] = const char Options::SCAN_FALLBACK_BRANCH[] = "scan.fallback-branch"; const char Options::BRANCH[] = "branch"; const char Options::FILE_INDEX_READ_ENABLED[] = "file-index.read.enabled"; +const char Options::FILE_INDEX_IN_MANIFEST_THRESHOLD[] = "file-index.in-manifest-threshold"; const char Options::DATA_FILE_EXTERNAL_PATHS[] = "data-file.external-paths"; const char Options::DATA_FILE_EXTERNAL_PATHS_STRATEGY[] = "data-file.external-paths.strategy"; const char Options::DATA_FILE_PREFIX[] = "data-file.prefix"; diff --git a/src/paimon/common/file_index/file_index_format.cpp b/src/paimon/common/file_index/file_index_format.cpp index 855008450..23bf3715c 100644 --- a/src/paimon/common/file_index/file_index_format.cpp +++ b/src/paimon/common/file_index/file_index_format.cpp @@ -27,7 +27,9 @@ #include "arrow/type.h" #include "fmt/format.h" #include "paimon/common/file_index/empty/empty_file_index_reader.h" +#include "paimon/common/io/data_output_stream.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/math.h" #include "paimon/file_index/file_indexer.h" #include "paimon/file_index/file_indexer_factory.h" #include "paimon/io/byte_array_input_stream.h" @@ -39,6 +41,102 @@ namespace paimon { class InputStream; class MemoryPool; +class FileIndexFormatWriterImpl : public FileIndexFormat::Writer { + public: + explicit FileIndexFormatWriterImpl(const std::shared_ptr& output_stream) + : output_stream_(output_stream) { + assert(output_stream_); + } + + Status WriteColumnIndexes(const FileIndexFormat::ColumnIndexes& indexes) override { + if (written_) { + return Status::Invalid("File index column indexes have already been written"); + } + int64_t header_length = sizeof(int64_t) + sizeof(int32_t) * 3 + sizeof(int32_t); + int64_t body_length = 0; + for (const auto& [column_name, column_indexes] : indexes) { + PAIMON_RETURN_NOT_OK(ValidateValueInRange(column_name.size(), + "file index column name length")); + header_length += + sizeof(uint16_t) + static_cast(column_name.size()) + sizeof(int32_t); + for (const auto& [index_type, bytes] : column_indexes) { + PAIMON_RETURN_NOT_OK(ValidateValueInRange(index_type.size(), + "file index type name length")); + header_length += sizeof(uint16_t) + static_cast(index_type.size()) + + sizeof(int32_t) * 2; + if (bytes) { + PAIMON_RETURN_NOT_OK(AddChecked(bytes->size(), &body_length, "index body")); + } + } + } + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(header_length, "file index header length")); + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(indexes.size(), "file index column count")); + PAIMON_RETURN_NOT_OK(AddChecked(header_length, &body_length, "file index size")); + + DataOutputStream data_output(output_stream_); + PAIMON_RETURN_NOT_OK(data_output.WriteValue(FileIndexFormat::MAGIC)); + PAIMON_RETURN_NOT_OK(data_output.WriteValue(FileIndexFormat::V_1)); + PAIMON_RETURN_NOT_OK(data_output.WriteValue(static_cast(header_length))); + PAIMON_RETURN_NOT_OK(data_output.WriteValue(static_cast(indexes.size()))); + + int64_t body_offset = header_length; + for (const auto& [column_name, column_indexes] : indexes) { + PAIMON_RETURN_NOT_OK(data_output.WriteString(column_name)); + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(column_indexes.size(), "column index count")); + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(static_cast(column_indexes.size()))); + for (const auto& [index_type, bytes] : column_indexes) { + PAIMON_RETURN_NOT_OK(data_output.WriteString(index_type)); + if (bytes == nullptr) { + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(FileIndexFormat::EMPTY_INDEX_FLAG)); + PAIMON_RETURN_NOT_OK(data_output.WriteValue(0)); + continue; + } + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(static_cast(body_offset))); + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(static_cast(bytes->size()))); + body_offset += static_cast(bytes->size()); + } + } + PAIMON_RETURN_NOT_OK(data_output.WriteValue(0)); + for (const auto& [column_name, column_indexes] : indexes) { + for (const auto& [index_type, bytes] : column_indexes) { + if (bytes) { + PAIMON_RETURN_NOT_OK(data_output.WriteBytes(bytes)); + } + } + } + written_ = true; + return Status::OK(); + } + + Status Close() override { + if (closed_) { + return Status::OK(); + } + closed_ = true; + PAIMON_RETURN_NOT_OK(output_stream_->Flush()); + return output_stream_->Close(); + } + + private: + template + static Status AddChecked(T value, int64_t* total, const char* name) { + PAIMON_RETURN_NOT_OK(ValidateValueInRange(value, name)); + *total += static_cast(value); + return ValidateValueInRange(*total, name); + } + + std::shared_ptr output_stream_; + bool written_ = false; + bool closed_ = false; +}; + class FileIndexFormatReaderImpl : public FileIndexFormat::Reader { public: using HeaderType = @@ -153,4 +251,15 @@ Result> FileIndexFormat::CreateReader( const std::shared_ptr& input_stream, const std::shared_ptr& pool) { return FileIndexFormatReaderImpl::Create(input_stream, pool); } + +Result> FileIndexFormat::CreateWriter( + const std::shared_ptr& output_stream, const std::shared_ptr& pool) { + if (!output_stream) { + return Status::Invalid("File index output stream cannot be null"); + } + if (!pool) { + return Status::Invalid("File index memory pool cannot be null"); + } + return std::make_unique(output_stream); +} } // namespace paimon diff --git a/src/paimon/common/file_index/file_index_format_test.cpp b/src/paimon/common/file_index/file_index_format_test.cpp index 7851d57e1..148c71b3d 100644 --- a/src/paimon/common/file_index/file_index_format_test.cpp +++ b/src/paimon/common/file_index/file_index_format_test.cpp @@ -24,17 +24,20 @@ #include "paimon/common/file_index/bloomfilter/bloom_filter_file_index.h" #include "paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.h" #include "paimon/common/file_index/empty/empty_file_index_reader.h" +#include "paimon/common/io/byte_array_output_stream.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" #include "paimon/file_index/file_index_result.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/io/byte_array_input_stream.h" +#include "paimon/memory/bytes.h" #include "paimon/memory/memory_pool.h" #include "paimon/predicate/literal.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { + class FileIndexFormatTest : public ::testing::Test { public: void SetUp() override { @@ -55,6 +58,24 @@ class FileIndexFormatTest : public ::testing::Test { std::shared_ptr pool_; }; +TEST_F(FileIndexFormatTest, TestWriteEmptyIndexGoldenBytes) { + // the expected bytes are generated from Java Paimon + std::vector expected = {0, 5, 78, 78, -48, 26, 53, -82, 0, 0, 0, 1, 0, 0, 0, 47, + 0, 0, 0, 1, 0, 2, 99, 49, 0, 0, 0, 1, 0, 5, 101, 109, + 112, 116, 121, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0}; + FileIndexFormat::ColumnIndexes indexes; + indexes["c1"]["empty"] = nullptr; + auto output = std::make_shared( + MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_); + + ASSERT_OK_AND_ASSIGN(auto writer, FileIndexFormat::CreateWriter(output, pool_)); + ASSERT_OK(writer->WriteColumnIndexes(indexes)); + ASSERT_OK(writer->Close()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, output->Finish()); + + ASSERT_EQ(expected, std::vector(actual->data(), actual->data() + actual->size())); +} + TEST_F(FileIndexFormatTest, TestCreateEmptyFileIndexReader) { auto schema = arrow::schema({arrow::field("c1", arrow::utf8())}); std::vector index_file_bytes = {0, 5, 78, 78, -48, 26, 53, -82, 0, 0, 0, 1, diff --git a/src/paimon/common/io/byte_array_output_stream.cpp b/src/paimon/common/io/byte_array_output_stream.cpp new file mode 100644 index 000000000..94da8deae --- /dev/null +++ b/src/paimon/common/io/byte_array_output_stream.cpp @@ -0,0 +1,79 @@ +/* + * 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/common/io/byte_array_output_stream.h" + +#include +#include +#include + +#include "paimon/common/memory/memory_segment_utils.h" +#include "paimon/common/utils/math.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" + +namespace paimon { + +ByteArrayOutputStream::ByteArrayOutputStream(int32_t initial_capacity, + const std::shared_ptr& pool) + : pool_(pool), output_(initial_capacity, pool_) {} + +Result ByteArrayOutputStream::Write(const char* buffer, int64_t size) { + if (closed_) { + return Status::Invalid("Byte array output stream is closed"); + } + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(size, "write length")); + if (buffer == nullptr && size > 0) { + return Status::Invalid("Write buffer must not be null when size is positive"); + } + int64_t remaining = size; + while (remaining > 0) { + uint32_t to_write = static_cast(std::min( + remaining, static_cast(std::numeric_limits::max()))); + output_.Write(buffer, to_write); + buffer += to_write; + remaining -= to_write; + } + position_ += size; + return size; +} + +Status ByteArrayOutputStream::Close() { + closed_ = true; + return Status::OK(); +} + +Result> ByteArrayOutputStream::Finish() { + PAIMON_RETURN_NOT_OK(Close()); + if (result_) { + return result_; + } + // TODO(jinli.zjw): Support int64_t lengths in MemorySegmentUtils::CopyToBytes and remove this + // limit. + if (position_ > std::numeric_limits::max()) { + return Status::Invalid("Byte array output stream size exceeds INT32_MAX"); + } + const std::vector& segments = output_.Segments(); + result_ = std::shared_ptr(new Bytes(static_cast(position_), pool_.get()), + [pool = pool_](Bytes* bytes) { delete bytes; }); + MemorySegmentUtils::CopyToBytes(segments, /*offset=*/0, result_.get(), + /*bytes_offset=*/0, static_cast(position_)); + return result_; +} + +} // namespace paimon diff --git a/src/paimon/common/io/byte_array_output_stream.h b/src/paimon/common/io/byte_array_output_stream.h new file mode 100644 index 000000000..bceec378d --- /dev/null +++ b/src/paimon/common/io/byte_array_output_stream.h @@ -0,0 +1,69 @@ +/* + * 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 +#include +#include + +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/fs/file_system.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace paimon { + +class Bytes; +class MemoryPool; + +/// An in-memory output stream backed by segments allocated from a Paimon MemoryPool. +class ByteArrayOutputStream : public OutputStream { + public: + ByteArrayOutputStream(int32_t initial_capacity, const std::shared_ptr& pool); + + ~ByteArrayOutputStream() override = default; + + Result Write(const char* buffer, int64_t size) override; + + Status Flush() override { + return Status::OK(); + } + + Result GetPos() const override { + return position_; + } + + Result GetUri() const override { + return std::string(); + } + + Status Close() override; + + /// Closes the stream and returns its contents as an exactly-sized contiguous byte array. + Result> Finish(); + + private: + std::shared_ptr pool_; + MemorySegmentOutputStream output_; + std::shared_ptr result_; + int64_t position_ = 0; + bool closed_ = false; +}; + +} // namespace paimon diff --git a/src/paimon/common/io/byte_array_output_stream_test.cpp b/src/paimon/common/io/byte_array_output_stream_test.cpp new file mode 100644 index 000000000..23ce567aa --- /dev/null +++ b/src/paimon/common/io/byte_array_output_stream_test.cpp @@ -0,0 +1,83 @@ +/* + * 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/common/io/byte_array_output_stream.h" + +#include +#include + +#include "gtest/gtest.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +TEST(ByteArrayOutputStreamTest, TestWriteAndFinish) { + std::shared_ptr pool = GetMemoryPool(); + std::shared_ptr stream = + std::make_shared(/*initial_capacity=*/2, pool); + ASSERT_GT(pool->CurrentUsage(), 0); + ASSERT_OK_AND_ASSIGN(int64_t first_write, stream->Write("ab", 2)); + ASSERT_EQ(2, first_write); + ASSERT_OK_AND_ASSIGN(int64_t second_write, stream->Write("cdef", 4)); + ASSERT_EQ(4, second_write); + ASSERT_OK_AND_ASSIGN(int64_t position, stream->GetPos()); + ASSERT_EQ(6, position); + ASSERT_EQ(pool->CurrentUsage(), pool->MaxMemoryUsage()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish()); + ASSERT_EQ("abcdef", std::string(result->data(), result->size())); + ASSERT_NOK_WITH_MSG(stream->Write("x", 1), "closed"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr repeated, stream->Finish()); + ASSERT_EQ(result, repeated); + stream.reset(); + ASSERT_EQ(6, pool->CurrentUsage()); +} + +TEST(ByteArrayOutputStreamTest, TestWriteValidation) { + std::shared_ptr stream = std::make_shared( + /*initial_capacity=*/8, GetDefaultPool()); + ASSERT_NOK(stream->Write(nullptr, 1)); + ASSERT_NOK(stream->Write("", -1)); + ASSERT_OK_AND_ASSIGN(int64_t written, stream->Write(nullptr, 0)); + ASSERT_EQ(0, written); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish()); + ASSERT_EQ(0, result->size()); +} + +TEST(ByteArrayOutputStreamTest, TestResultKeepsMemoryPoolAlive) { + std::shared_ptr pool = GetMemoryPool(); + std::weak_ptr weak_pool = pool; + std::shared_ptr stream = + std::make_shared(/*initial_capacity=*/8, pool); + ASSERT_OK_AND_ASSIGN(int64_t written, stream->Write("data", 4)); + ASSERT_EQ(4, written); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish()); + + stream.reset(); + pool.reset(); + ASSERT_FALSE(weak_pool.expired()); + ASSERT_EQ("data", std::string(result->data(), result->size())); + + result.reset(); + ASSERT_TRUE(weak_pool.expired()); +} + +} // namespace paimon::test diff --git a/src/paimon/common/io/memory_segment_output_stream.cpp b/src/paimon/common/io/memory_segment_output_stream.cpp index 5355f72ba..cbcfabf9c 100644 --- a/src/paimon/common/io/memory_segment_output_stream.cpp +++ b/src/paimon/common/io/memory_segment_output_stream.cpp @@ -54,12 +54,20 @@ void MemorySegmentOutputStream::WriteString(const std::string& str) { } void MemorySegmentOutputStream::Write(const char* data, uint32_t size) { - auto bytes = std::make_shared(size, pool_.get()); - if (size != 0) { - memcpy(bytes->data(), data, size); + uint32_t remaining_size = size; + while (remaining_size > 0) { + int32_t remaining_in_segment = segment_size_ - position_in_segment_; + if (remaining_in_segment == 0) { + Advance(); + remaining_in_segment = segment_size_; + } + int32_t to_write = static_cast( + std::min(remaining_size, static_cast(remaining_in_segment))); + std::memcpy(current_segment_.MutableData() + position_in_segment_, data, to_write); + data += to_write; + remaining_size -= to_write; + position_in_segment_ += to_write; } - auto segment = MemorySegment::Wrap(bytes); - Write(segment, 0, segment.Size()); } void MemorySegmentOutputStream::Write(const MemorySegment& segment, int32_t offset, int32_t len) { diff --git a/src/paimon/common/io/memory_segment_output_stream_test.cpp b/src/paimon/common/io/memory_segment_output_stream_test.cpp index 61fbfe305..69c207ce9 100644 --- a/src/paimon/common/io/memory_segment_output_stream_test.cpp +++ b/src/paimon/common/io/memory_segment_output_stream_test.cpp @@ -82,4 +82,16 @@ TEST_P(MemorySegmentOutputStreamTest, TestSimple) { ASSERT_EQ(out.CurrentSize(), input_stream->GetPos().value()); } +TEST(MemorySegmentOutputStreamTest, TestRawWriteDoesNotAllocateTemporaryBuffer) { + std::shared_ptr pool = GetMemoryPool(); + MemorySegmentOutputStream out(/*segment_size=*/8, pool); + uint64_t allocated_before_write = pool->CurrentUsage(); + + out.Write("abc", 3); + + ASSERT_EQ(allocated_before_write, pool->CurrentUsage()); + ASSERT_EQ(pool->CurrentUsage(), pool->MaxMemoryUsage()); + ASSERT_EQ(3, out.CurrentSize()); +} + } // namespace paimon::test diff --git a/src/paimon/core/append/append_only_writer_test.cpp b/src/paimon/core/append/append_only_writer_test.cpp index 8ef6e135f..114016091 100644 --- a/src/paimon/core/append/append_only_writer_test.cpp +++ b/src/paimon/core/append/append_only_writer_test.cpp @@ -56,9 +56,11 @@ #include "paimon/core/stats/simple_stats.h" #include "paimon/core/utils/commit_increment.h" #include "paimon/defs.h" +#include "paimon/file_index/file_index_format.h" #include "paimon/format/file_format_factory.h" #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/io/byte_array_input_stream.h" #include "paimon/memory/memory_pool.h" #include "paimon/record_batch.h" #include "paimon/testing/utils/binary_row_generator.h" @@ -404,6 +406,41 @@ TEST_F(AppendOnlyWriterTest, TestWriteAndPrepareCommit) { ASSERT_OK(writer->Close()); } +TEST_F(AppendOnlyWriterTest, TestWritePublishesEmbeddedBitmapIndex) { + CoreOptions options = CreateOptions( + {{"file-index.bitmap.columns", "f0"}, {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}}); + auto schema = + arrow::schema({arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::int32())}); + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = CreatePathFactory(dir->Str(), "mock_format", options); + ASSERT_OK_AND_ASSIGN( + auto writer, CreateAppendOnlyWriter( + options, /*schema_id=*/0, schema, /*write_cols=*/std::nullopt, + /*max_sequence_number=*/-1, path_factory, compact_manager_, memory_pool_)); + + ASSERT_OK(writer->Write(CreateBatch(schema, R"([{"f0": 1, "f1": 10}, + {"f0": 2, "f1": 20}, + {"f0": 1, "f1": 30}])"))); + ASSERT_OK_AND_ASSIGN(CommitIncrement increment, + writer->PrepareCommit(/*wait_compaction=*/true)); + const auto& files = increment.GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + ASSERT_TRUE(files[0]->embedded_index); + ASSERT_TRUE(files[0]->extra_files.empty()); + + auto input = std::make_shared(files[0]->embedded_index->data(), + files[0]->embedded_index->size()); + ASSERT_OK_AND_ASSIGN(auto index_reader, FileIndexFormat::CreateReader(input, memory_pool_)); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); + ASSERT_OK_AND_ASSIGN(auto column_readers, index_reader->ReadColumnIndex("f0", &c_schema)); + ASSERT_EQ(1, column_readers.size()); + ASSERT_OK_AND_ASSIGN(auto result, column_readers[0]->VisitEqual(Literal(1))); + ASSERT_EQ("{0,2}", result->ToString()); + ASSERT_OK(writer->Close()); +} + TEST_F(AppendOnlyWriterTest, TestWriteAndClose) { std::map raw_options; raw_options[Options::FILE_FORMAT] = "orc"; diff --git a/src/paimon/core/core_options.cpp b/src/paimon/core/core_options.cpp index 578950fbe..1c2e164bf 100644 --- a/src/paimon/core/core_options.cpp +++ b/src/paimon/core/core_options.cpp @@ -386,6 +386,7 @@ struct CoreOptions::Impl { int64_t manifest_target_file_size = 8 * 1024 * 1024; int64_t deletion_vector_target_file_size = 2 * 1024 * 1024; int64_t manifest_full_compaction_file_size = 16 * 1024 * 1024; + int64_t file_index_in_manifest_threshold = 500; int64_t write_buffer_size = 256 * 1024 * 1024; int64_t commit_timeout = std::numeric_limits::max(); int64_t commit_min_retry_wait = 10; @@ -838,6 +839,9 @@ struct CoreOptions::Impl { // Parse index-related configurations: file index, global index. Status ParseIndexOptions(const ConfigParser& parser) { + // Parse file-index.in-manifest-threshold - max inline file index size, default 500B + PAIMON_RETURN_NOT_OK(parser.ParseMemorySize(Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, + &file_index_in_manifest_threshold)); // Parse file-index.read.enabled - whether to enable reading file index, default true PAIMON_RETURN_NOT_OK( parser.Parse(Options::FILE_INDEX_READ_ENABLED, &file_index_read_enabled)); @@ -1654,6 +1658,10 @@ bool CoreOptions::FileIndexReadEnabled() const { return impl_->file_index_read_enabled; } +int64_t CoreOptions::FileIndexInManifestThreshold() const { + return impl_->file_index_in_manifest_threshold; +} + std::optional CoreOptions::GetDataFileExternalPaths() const { return impl_->data_file_external_paths; } diff --git a/src/paimon/core/core_options.h b/src/paimon/core/core_options.h index 9eb289883..53ef4ad0f 100644 --- a/src/paimon/core/core_options.h +++ b/src/paimon/core/core_options.h @@ -207,6 +207,7 @@ class PAIMON_EXPORT CoreOptions { bool NeedLookup() const; bool PrepareCommitWaitCompaction() const; bool FileIndexReadEnabled() const; + int64_t FileIndexInManifestThreshold() const; std::map GetFieldsSequenceGroups() const; bool PartialUpdateRemoveRecordOnDelete() const; diff --git a/src/paimon/core/core_options_test.cpp b/src/paimon/core/core_options_test.cpp index ebc127edb..0054a5b5f 100644 --- a/src/paimon/core/core_options_test.cpp +++ b/src/paimon/core/core_options_test.cpp @@ -134,6 +134,7 @@ TEST(CoreOptionsTest, TestDefaultValue) { ASSERT_EQ(std::nullopt, core_options.GetScanFallbackBranch()); ASSERT_EQ("main", core_options.GetBranch()); ASSERT_TRUE(core_options.FileIndexReadEnabled()); + ASSERT_EQ(500, core_options.FileIndexInManifestThreshold()); ASSERT_EQ(std::nullopt, core_options.GetDataFileExternalPaths()); ASSERT_EQ(ExternalPathStrategy::NONE, core_options.GetExternalPathStrategy()); ASSERT_TRUE(core_options.EnableAdaptivePrefetchStrategy()); @@ -248,6 +249,7 @@ TEST(CoreOptionsTest, TestFromMap) { {Options::SCAN_FALLBACK_BRANCH, "fallback"}, {Options::BRANCH, "rt"}, {Options::FILE_INDEX_READ_ENABLED, "false"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "2KB"}, {Options::DATA_FILE_EXTERNAL_PATHS, "FILE:///tmp/index"}, {Options::DATA_FILE_EXTERNAL_PATHS_STRATEGY, "round-robin"}, {Options::FILE_COMPRESSION, "snappy"}, @@ -398,6 +400,7 @@ TEST(CoreOptionsTest, TestFromMap) { ASSERT_EQ(core_options.GetScanFallbackBranch(), std::optional("fallback")); ASSERT_EQ(core_options.GetBranch(), "rt"); ASSERT_FALSE(core_options.FileIndexReadEnabled()); + ASSERT_EQ(2 * 1024, core_options.FileIndexInManifestThreshold()); ASSERT_EQ(core_options.GetDataFileExternalPaths(), std::optional("FILE:///tmp/index")); ASSERT_EQ(core_options.GetExternalPathStrategy(), ExternalPathStrategy::ROUND_ROBIN); diff --git a/src/paimon/core/io/append_data_file_writer_factory.cpp b/src/paimon/core/io/append_data_file_writer_factory.cpp index e12374ddb..d0b677a78 100644 --- a/src/paimon/core/io/append_data_file_writer_factory.cpp +++ b/src/paimon/core/io/append_data_file_writer_factory.cpp @@ -24,6 +24,7 @@ #include "arrow/c/abi.h" #include "arrow/c/helpers.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/fs/file_system.h" @@ -57,6 +58,11 @@ AppendDataFileWriterFactory::CreateWriter() const { options_.GetFileCompression(), std::function(), schema_id_, seq_num_counter, file_source_, resources.stats_extractor, path_factory_->IsExternalPath(), write_cols_, pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_index_writer, + CreateFileIndexWriter(write_schema_, path_factory_)); + if (file_index_writer) { + writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); return std::unique_ptr>>( diff --git a/src/paimon/core/io/data_file_index_writer.cpp b/src/paimon/core/io/data_file_index_writer.cpp new file mode 100644 index 000000000..844a30e4a --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -0,0 +1,164 @@ +/* + * 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/io/data_file_index_writer.h" + +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/io/byte_array_output_stream.h" +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/io/data_file_path_factory.h" +#include "paimon/file_index/file_index_format.h" +#include "paimon/file_index/file_index_writer.h" +#include "paimon/file_index/file_indexer.h" +#include "paimon/file_index/file_indexer_factory.h" +#include "paimon/fs/file_system.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/status.h" + +namespace paimon { +Result> DataFileIndexWriter::Create( + const std::shared_ptr& logical_schema, const FileIndexOptions& options, + const std::shared_ptr& file_system, + const std::shared_ptr& path_factory, + const std::shared_ptr& pool) { + assert(logical_schema); + assert(file_system); + assert(path_factory); + assert(pool); + std::vector writers; + writers.reserve(options.Definitions().size()); + for (const FileIndexDefinition& definition : options.Definitions()) { + int32_t field_index = logical_schema->GetFieldIndex(definition.column_name); + if (field_index < 0) { + return Status::Invalid( + fmt::format("File index column '{}' does not exist in the write schema", + definition.column_name)); + } + std::shared_ptr field = logical_schema->field(field_index); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr indexer, + FileIndexerFactory::Get(definition.index_type, definition.options)); + if (!indexer) { + return Status::Invalid( + fmt::format("File index type '{}' is not registered", definition.index_type)); + } + ::ArrowSchema c_schema; + ArrowSchemaMarkReleased(&c_schema); + ScopeGuard schema_guard([&c_schema]() { ArrowSchemaRelease(&c_schema); }); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow::schema({field}), &c_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr writer, + indexer->CreateWriter(&c_schema, pool)); + writers.push_back( + {definition.column_name, definition.index_type, field_index, field, std::move(writer)}); + } + return std::unique_ptr(new DataFileIndexWriter( + std::move(writers), options.InManifestThreshold(), file_system, path_factory, pool)); +} + +DataFileIndexWriter::DataFileIndexWriter(std::vector&& writers, + int64_t in_manifest_threshold, + const std::shared_ptr& file_system, + const std::shared_ptr& path_factory, + const std::shared_ptr& pool) + : writers_(std::move(writers)), + in_manifest_threshold_(in_manifest_threshold), + file_system_(file_system), + path_factory_(path_factory), + pool_(pool) {} + +Status DataFileIndexWriter::AddBatch(const std::shared_ptr& logical_batch) { + for (IndexWriterEntry& entry : writers_) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr projected, + arrow::StructArray::Make({logical_batch->field(entry.field_index)}, {entry.field})); + ::ArrowArray c_array; + ArrowArrayMarkReleased(&c_array); + ScopeGuard array_guard([&c_array]() { ArrowArrayRelease(&c_array); }); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*projected, &c_array)); + PAIMON_RETURN_NOT_OK(entry.writer->AddBatch(&c_array)); + } + return Status::OK(); +} + +Result> DataFileIndexWriter::SerializeContainer() { + FileIndexFormat::ColumnIndexes column_indexes; + for (IndexWriterEntry& entry : writers_) { + PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR serialized, + entry.writer->SerializedBytes()); + column_indexes[entry.column_name][entry.index_type] = + std::shared_ptr(std::move(serialized)); + } + + std::shared_ptr output = std::make_shared( + MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_writer, + FileIndexFormat::CreateWriter(output, pool_)); + PAIMON_RETURN_NOT_OK(format_writer->WriteColumnIndexes(column_indexes)); + PAIMON_RETURN_NOT_OK(format_writer->Close()); + return output->Finish(); +} + +Result DataFileIndexWriter::Finish(const std::string& data_file_path) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, SerializeContainer()); + if (static_cast(bytes->size()) <= in_manifest_threshold_) { + return FileIndexWriteResult{bytes, {}}; + } + + external_index_path_ = path_factory_->ToFileIndexPath(data_file_path); + PAIMON_RETURN_NOT_OK(WriteExternal(external_index_path_.value(), bytes)); + return FileIndexWriteResult{nullptr, {PathUtil::GetName(external_index_path_.value())}}; +} + +Status DataFileIndexWriter::WriteExternal(const std::string& path, + const std::shared_ptr& bytes) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, + file_system_->Create(path, /*overwrite=*/false)); + ScopeGuard guard([this, &output]() { + [[maybe_unused]] Status _ = output->Close(); + Abort(); + }); + PAIMON_ASSIGN_OR_RAISE(int64_t written, + output->Write(bytes->data(), static_cast(bytes->size()))); + if (written != static_cast(bytes->size())) { + return Status::IOError(fmt::format("Short write for file index {}: expected {}, wrote {}", + path, bytes->size(), written)); + } + PAIMON_RETURN_NOT_OK(output->Flush()); + PAIMON_RETURN_NOT_OK(output->Close()); + output.reset(); + guard.Release(); + return Status::OK(); +} + +void DataFileIndexWriter::Abort() { + if (external_index_path_) { + [[maybe_unused]] Status _ = file_system_->Delete(external_index_path_.value()); + } +} + +} // namespace paimon diff --git a/src/paimon/core/io/data_file_index_writer.h b/src/paimon/core/io/data_file_index_writer.h new file mode 100644 index 000000000..ffb17511b --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer.h @@ -0,0 +1,93 @@ +/* + * 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 +#include +#include +#include + +#include "paimon/core/io/file_index_options.h" +#include "paimon/result.h" + +namespace arrow { +class Field; +class Schema; +class StructArray; +} // namespace arrow + +namespace paimon { + +class Bytes; +class DataFilePathFactory; +class FileIndexWriter; +class FileSystem; +class MemoryPool; + +struct FileIndexWriteResult { + std::shared_ptr embedded_index; + std::vector> extra_files; +}; + +/// Builds every configured column index for one data file. +class DataFileIndexWriter { + public: + static Result> Create( + const std::shared_ptr& logical_schema, const FileIndexOptions& options, + const std::shared_ptr& file_system, + const std::shared_ptr& path_factory, + const std::shared_ptr& pool); + + Status AddBatch(const std::shared_ptr& logical_batch); + + Result Finish(const std::string& data_file_path); + + void Abort(); + + const std::optional& ExternalIndexPath() const { + return external_index_path_; + } + + private: + struct IndexWriterEntry { + std::string column_name; + std::string index_type; + int32_t field_index; + std::shared_ptr field; + std::shared_ptr writer; + }; + + DataFileIndexWriter(std::vector&& writers, int64_t in_manifest_threshold, + const std::shared_ptr& file_system, + const std::shared_ptr& path_factory, + const std::shared_ptr& pool); + + Result> SerializeContainer(); + Status WriteExternal(const std::string& path, const std::shared_ptr& bytes); + + std::vector writers_; + int64_t in_manifest_threshold_; + std::shared_ptr file_system_; + std::shared_ptr path_factory_; + std::shared_ptr pool_; + std::optional external_index_path_; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/data_file_index_writer_test.cpp b/src/paimon/core/io/data_file_index_writer_test.cpp new file mode 100644 index 000000000..5e881b4e1 --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -0,0 +1,161 @@ +/* + * 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/io/data_file_index_writer.h" + +#include +#include +#include + +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "arrow/type.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_path_factory.h" +#include "paimon/core/io/file_index_options.h" +#include "paimon/defs.h" +#include "paimon/file_index/file_index_format.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/io/byte_array_input_stream.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class DataFileIndexWriterTest : public ::testing::Test { + public: + void SetUp() override { + pool_ = GetDefaultPool(); + file_system_ = std::make_shared(); + directory_ = UniqueTestDirectory::Create(); + ASSERT_TRUE(directory_); + path_factory_ = std::make_shared(); + ASSERT_OK(path_factory_->Init(directory_->Str(), "orc", "data-", nullptr)); + schema_ = + arrow::schema({arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::int32())}); + } + + Result> CreateWriter( + const std::map& index_options) const { + std::map options = {{"file-system", "local"}}; + options.insert(index_options.begin(), index_options.end()); + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, + CoreOptions::FromMap(options, file_system_)); + PAIMON_ASSIGN_OR_RAISE(FileIndexOptions parsed, + FileIndexOptions::FromCoreOptions(core_options)); + return DataFileIndexWriter::Create(schema_, parsed, file_system_, path_factory_, pool_); + } + + std::shared_ptr CreateBatch(const std::string& json) const { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema_->fields()), json) + .ValueOrDie(); + return std::dynamic_pointer_cast(array); + } + + Result> CreateReader( + const std::shared_ptr& bytes) const { + auto input = std::make_shared(bytes->data(), bytes->size()); + return FileIndexFormat::CreateReader(input, pool_); + } + + Result>> ReadColumn( + FileIndexFormat::Reader* reader, const std::string& column_name) const { + ::ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, &c_schema)); + return reader->ReadColumnIndex(column_name, &c_schema); + } + + protected: + std::shared_ptr pool_; + std::shared_ptr file_system_; + std::unique_ptr directory_; + std::shared_ptr path_factory_; + std::shared_ptr schema_; +}; + +TEST_F(DataFileIndexWriterTest, TestBitmapAndRangeBitmapEmbeddedRoundTrip) { + ASSERT_OK_AND_ASSIGN(auto writer, + CreateWriter({{"file-index.bitmap.columns", "f0"}, + {"file-index.range-bitmap.columns", "f1"}, + {"file-index.range-bitmap.f1.chunk-size", "1KB"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}})); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 10}, + {"f0": 2, "f1": 20}])"))); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 30}, + {"f0": null, "f1": 40}])"))); + + ASSERT_OK_AND_ASSIGN(FileIndexWriteResult result, writer->Finish("unused.orc")); + ASSERT_TRUE(result.embedded_index); + ASSERT_TRUE(result.extra_files.empty()); + ASSERT_OK_AND_ASSIGN(auto reader, CreateReader(result.embedded_index)); + + ASSERT_OK_AND_ASSIGN(auto bitmap_readers, ReadColumn(reader.get(), "f0")); + ASSERT_EQ(1, bitmap_readers.size()); + ASSERT_OK_AND_ASSIGN(auto equal_result, bitmap_readers[0]->VisitEqual(Literal(1))); + ASSERT_EQ("{0,2}", equal_result->ToString()); + ASSERT_OK_AND_ASSIGN(auto null_result, bitmap_readers[0]->VisitIsNull()); + ASSERT_EQ("{3}", null_result->ToString()); + + ASSERT_OK_AND_ASSIGN(auto range_readers, ReadColumn(reader.get(), "f1")); + ASSERT_EQ(1, range_readers.size()); + ASSERT_OK_AND_ASSIGN(auto greater_result, range_readers[0]->VisitGreaterThan(Literal(20))); + ASSERT_EQ("{2,3}", greater_result->ToString()); +} + +TEST_F(DataFileIndexWriterTest, TestExternalIndexAndAbortCleanup) { + ASSERT_OK_AND_ASSIGN(auto writer, + CreateWriter({{"file-index.bitmap.columns", "f0"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1B"}})); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 10}])"))); + std::string data_path = path_factory_->NewPath(); + + ASSERT_OK_AND_ASSIGN(FileIndexWriteResult result, writer->Finish(data_path)); + ASSERT_FALSE(result.embedded_index); + ASSERT_EQ(1, result.extra_files.size()); + ASSERT_TRUE(result.extra_files[0]); + ASSERT_EQ(PathUtil::GetName(path_factory_->ToFileIndexPath(data_path)), + result.extra_files[0].value()); + std::string index_path = path_factory_->ToFileIndexPath(data_path); + ASSERT_OK_AND_ASSIGN(bool exists, file_system_->Exists(index_path)); + ASSERT_TRUE(exists); + ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system_->Open(index_path)); + ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input, pool_)); + ASSERT_OK_AND_ASSIGN(auto bitmap_readers, ReadColumn(reader.get(), "f0")); + ASSERT_EQ(1, bitmap_readers.size()); + ASSERT_OK_AND_ASSIGN(auto equal_result, bitmap_readers[0]->VisitEqual(Literal(1))); + ASSERT_EQ("{0}", equal_result->ToString()); + + writer->Abort(); + ASSERT_OK_AND_ASSIGN(exists, file_system_->Exists(index_path)); + ASSERT_FALSE(exists); +} + +TEST_F(DataFileIndexWriterTest, TestUnavailableWriterFailsCreation) { + ASSERT_NOK_WITH_MSG(CreateWriter({{"file-index.unknown.columns", "f0"}}), + "File index type 'unknown' is not registered"); + ASSERT_NOK_WITH_MSG(CreateWriter({{"file-index.bloom-filter.columns", "f0"}}), + "do not support index writer in bloom filter"); +} + +} // namespace paimon::test diff --git a/src/paimon/core/io/data_file_writer.cpp b/src/paimon/core/io/data_file_writer.cpp index 4ed3e040c..5c3f4c2fe 100644 --- a/src/paimon/core/io/data_file_writer.cpp +++ b/src/paimon/core/io/data_file_writer.cpp @@ -19,6 +19,7 @@ #include "paimon/core/io/data_file_writer.h" #include +#include #include "arrow/c/abi.h" #include "paimon/common/utils/long_counter.h" @@ -36,7 +37,7 @@ DataFileWriter::DataFileWriter( const std::shared_ptr& stats_extractor, bool is_external_path, const std::optional>& write_cols, const std::shared_ptr& pool) - : SingleFileWriter(compression, converter), + : DataFileWriterBase(compression, std::move(converter)), pool_(pool), schema_id_(schema_id), is_external_path_(is_external_path), @@ -45,28 +46,13 @@ DataFileWriter::DataFileWriter( stats_extractor_(stats_extractor), write_cols_(write_cols) {} -void DataFileWriter::SetMetadataFinalizer(MetadataFinalizer finalizer) { - metadata_finalizer_ = std::move(finalizer); -} - Status DataFileWriter::Write(ArrowArray* batch) { int64_t record_count = batch->length; - PAIMON_RETURN_NOT_OK(SingleFileWriter::Write(batch)); + PAIMON_RETURN_NOT_OK(WriteRecord(batch, batch)); seq_num_counter_->Add(record_count); return Status::OK(); } -Status DataFileWriter::BeforeFinish() { - if (metadata_finalizer_) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr updated_schema, - metadata_finalizer_()); - if (updated_schema) { - PAIMON_RETURN_NOT_OK(UpdateSchema(updated_schema)); - } - } - return Status::OK(); -} - Result> DataFileWriter::GetResult() { PAIMON_ASSIGN_OR_RAISE(std::vector> field_stats, GetFieldStats()); PAIMON_ASSIGN_OR_RAISE(SimpleStats stats, @@ -77,11 +63,12 @@ Result> DataFileWriter::GetResult() { PAIMON_ASSIGN_OR_RAISE(Path external_path, PathUtil::ToPath(path_)); final_path = external_path.ToString(); } + const FileIndexWriteResult& file_index = GetFileIndexWriteResult(); return DataFileMeta::ForAppend( PathUtil::GetName(path_), output_bytes_, RecordCount(), stats, seq_num_counter_->GetValue() - RecordCount(), seq_num_counter_->GetValue() - 1, schema_id_, - {}, /*embedded_index=*/nullptr, file_source_, /*value_stats_cols=*/std::nullopt, final_path, - /*first_row_id=*/std::nullopt, write_cols_); + file_index.extra_files, file_index.embedded_index, file_source_, + /*value_stats_cols=*/std::nullopt, final_path, /*first_row_id=*/std::nullopt, write_cols_); } Result>> DataFileWriter::GetFieldStats() { diff --git a/src/paimon/core/io/data_file_writer.h b/src/paimon/core/io/data_file_writer.h index 60cc808a9..f56f34956 100644 --- a/src/paimon/core/io/data_file_writer.h +++ b/src/paimon/core/io/data_file_writer.h @@ -28,7 +28,7 @@ #include "arrow/c/abi.h" #include "paimon/common/utils/long_counter.h" #include "paimon/core/io/data_file_meta.h" -#include "paimon/core/io/single_file_writer.h" +#include "paimon/core/io/data_file_writer_base.h" #include "paimon/core/manifest/file_source.h" #include "paimon/result.h" #include "paimon/status.h" @@ -44,13 +44,8 @@ class FormatStatsExtractor; class LongCounter; class MemoryPool; -class DataFileWriter : public SingleFileWriter<::ArrowArray*, std::shared_ptr> { +class DataFileWriter : public DataFileWriterBase<::ArrowArray*> { public: - /// Callback invoked during BeforeFinish() to finalize file metadata. - /// Produces an updated schema with per-field metadata (e.g. shredding metadata) - /// and may perform other finalization work (e.g. reporting stats to cross-file context). - using MetadataFinalizer = std::function>()>; - DataFileWriter(const std::string& compression, std::function converter, int64_t schema_id, const std::shared_ptr& seq_num_counter, FileSource file_source, @@ -58,17 +53,10 @@ class DataFileWriter : public SingleFileWriter<::ArrowArray*, std::shared_ptr>& write_cols, const std::shared_ptr& pool); - /// Sets the metadata finalizer. Called during BeforeFinish() to produce an updated - /// schema and perform finalization callbacks. Must be set before Close(). - void SetMetadataFinalizer(MetadataFinalizer finalizer); - Status Write(::ArrowArray* batch) override; Result> GetResult() override; - protected: - Status BeforeFinish() override; - private: Result>> GetFieldStats(); @@ -80,7 +68,6 @@ class DataFileWriter : public SingleFileWriter<::ArrowArray*, std::shared_ptr stats_extractor_; std::optional> write_cols_; - MetadataFinalizer metadata_finalizer_; }; } // namespace paimon diff --git a/src/paimon/core/io/data_file_writer_base.h b/src/paimon/core/io/data_file_writer_base.h new file mode 100644 index 000000000..64b556857 --- /dev/null +++ b/src/paimon/core/io/data_file_writer_base.h @@ -0,0 +1,134 @@ +/* + * 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 +#include +#include +#include + +#include "arrow/c/bridge.h" +#include "arrow/type.h" +#include "arrow/util/checked_cast.h" +#include "paimon/core/io/data_file_index_writer.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/io/single_file_writer.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +/// Common lifecycle for data file writers which may finalize schema metadata and publish a +/// file-level index. Concrete writers remain responsible for their record-specific state and +/// DataFileMeta construction. +template +class DataFileWriterBase : public SingleFileWriter> { + public: + using Base = SingleFileWriter>; + using AbortExecutor = typename Base::AbortExecutor; + /// Callback invoked during BeforeFinish() to finalize file metadata. + /// Produces an updated schema with per-field metadata (e.g. shredding metadata) + /// and may perform other finalization work (e.g. reporting stats to cross-file context). + using MetadataFinalizer = std::function>()>; + + /// Sets the metadata finalizer. Called during BeforeFinish() to produce an updated + /// schema and perform finalization callbacks. Must be set before Close(). + void SetMetadataFinalizer(MetadataFinalizer finalizer) { + metadata_finalizer_ = std::move(finalizer); + } + + void SetFileIndexWriter(std::unique_ptr file_index_writer, + const std::shared_ptr& logical_schema) { + file_index_writer_ = std::move(file_index_writer); + logical_type_ = arrow::struct_(logical_schema->fields()); + } + + void Abort() override { + if (file_index_writer_) { + // The external index uses a path different from the data file path deleted by Base. + file_index_writer_->Abort(); + } + Base::Abort(); + } + + Result GetAbortExecutor() const override { + PAIMON_ASSIGN_OR_RAISE(AbortExecutor executor, Base::GetAbortExecutor()); + if (file_index_writer_ && file_index_writer_->ExternalIndexPath()) { + executor.Add(this->fs_, file_index_writer_->ExternalIndexPath().value()); + } + return executor; + } + + protected: + DataFileWriterBase(const std::string& compression, + std::function converter) + : Base(compression, std::move(converter)) {} + + Status WriteRecord(Record record, ::ArrowArray* logical_batch) { + PAIMON_RETURN_NOT_OK(AddFileIndexBatch(logical_batch)); + return Base::Write(std::move(record)); + } + + const FileIndexWriteResult& GetFileIndexWriteResult() const { + return file_index_result_; + } + + Status BeforeFinish() override { + if (metadata_finalizer_) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr updated_schema, + metadata_finalizer_()); + if (updated_schema) { + PAIMON_RETURN_NOT_OK(this->UpdateSchema(updated_schema)); + } + } + return Status::OK(); + } + + Status BeforeCompletion() override { + if (file_index_writer_) { + PAIMON_ASSIGN_OR_RAISE(file_index_result_, file_index_writer_->Finish(this->path_)); + } + return Status::OK(); + } + + private: + Status AddFileIndexBatch(::ArrowArray* batch) { + if (!file_index_writer_) { + return Status::OK(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr logical_array, + arrow::ImportArray(batch, logical_type_)); + std::shared_ptr logical_batch = + arrow::internal::checked_pointer_cast(logical_array); + PAIMON_RETURN_NOT_OK(file_index_writer_->AddBatch(logical_batch)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*logical_batch, batch)); + return Status::OK(); + } + + MetadataFinalizer metadata_finalizer_; + std::unique_ptr file_index_writer_; + std::shared_ptr logical_type_; + FileIndexWriteResult file_index_result_; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/data_file_writer_factory.cpp b/src/paimon/core/io/data_file_writer_factory.cpp index b929dde83..07195ab7f 100644 --- a/src/paimon/core/io/data_file_writer_factory.cpp +++ b/src/paimon/core/io/data_file_writer_factory.cpp @@ -24,6 +24,9 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" +#include "paimon/core/io/data_file_path_factory.h" +#include "paimon/core/io/file_index_options.h" #include "paimon/format/file_format.h" #include "paimon/format/writer_builder.h" @@ -58,4 +61,16 @@ Result DataFileWriterFactory::CreateWrit return resources; } +Result> DataFileWriterFactory::CreateFileIndexWriter( + const std::shared_ptr& logical_schema, + const std::shared_ptr& path_factory) const { + PAIMON_ASSIGN_OR_RAISE(FileIndexOptions file_index_options, + FileIndexOptions::FromCoreOptions(options_)); + if (file_index_options.Empty()) { + return std::unique_ptr(); + } + return DataFileIndexWriter::Create(logical_schema, file_index_options, options_.GetFileSystem(), + path_factory, pool_); +} + } // namespace paimon diff --git a/src/paimon/core/io/data_file_writer_factory.h b/src/paimon/core/io/data_file_writer_factory.h index c727b47d0..cab942f0e 100644 --- a/src/paimon/core/io/data_file_writer_factory.h +++ b/src/paimon/core/io/data_file_writer_factory.h @@ -32,6 +32,8 @@ class Schema; namespace paimon { class FileFormat; +class DataFileIndexWriter; +class DataFilePathFactory; class FormatStatsExtractor; class MemoryPool; class WriterBuilder; @@ -52,6 +54,10 @@ class DataFileWriterFactory { const std::shared_ptr& file_schema, bool create_stats_extractor) const; + Result> CreateFileIndexWriter( + const std::shared_ptr& logical_schema, + const std::shared_ptr& path_factory) const; + CoreOptions options_; int64_t schema_id_; std::shared_ptr pool_; diff --git a/src/paimon/core/io/file_index_options.cpp b/src/paimon/core/io/file_index_options.cpp new file mode 100644 index 000000000..0b170d612 --- /dev/null +++ b/src/paimon/core/io/file_index_options.cpp @@ -0,0 +1,102 @@ +/* + * 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/io/file_index_options.h" + +#include +#include + +#include "fmt/format.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/core/core_options.h" +#include "paimon/defs.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +constexpr char kFileIndexPrefix[] = "file-index."; +constexpr char kColumnsSuffix[] = ".columns"; + +} // namespace + +Result FileIndexOptions::FromCoreOptions(const CoreOptions& options) { + FileIndexOptions result; + const std::map& raw_options = options.ToMap(); + result.in_manifest_threshold_ = options.FileIndexInManifestThreshold(); + + std::set> declared; + for (const auto& [key, value] : raw_options) { + if (!StringUtils::StartsWith(key, kFileIndexPrefix) || + !StringUtils::EndsWith(key, kColumnsSuffix)) { + continue; + } + const size_t index_type_length = + key.size() - std::string(kFileIndexPrefix).size() - std::string(kColumnsSuffix).size(); + const std::string index_type = + key.substr(std::string(kFileIndexPrefix).size(), index_type_length); + if (index_type.empty()) { + return Status::Invalid(fmt::format("Invalid file index option {}", key)); + } + for (std::string column_name : StringUtils::Split(value, ",", /*ignore_empty=*/false)) { + StringUtils::Trim(&column_name); + if (column_name.empty()) { + return Status::Invalid( + fmt::format("Wrong option in {}, should not have empty column", key)); + } + if (column_name.find('[') != std::string::npos) { + return Status::NotImplemented( + "Writing file indexes for nested map columns is not supported"); + } + if (declared.emplace(column_name, index_type).second) { + result.definitions_.push_back({column_name, index_type, {}}); + } + } + } + + for (const auto& [key, value] : raw_options) { + if (!StringUtils::StartsWith(key, kFileIndexPrefix) || + StringUtils::EndsWith(key, kColumnsSuffix) || + key == Options::FILE_INDEX_IN_MANIFEST_THRESHOLD) { + continue; + } + std::vector parts = StringUtils::Split( + key.substr(std::string(kFileIndexPrefix).size()), ".", /*ignore_empty=*/false); + if (parts.size() != 3) { + continue; + } + bool found = false; + for (FileIndexDefinition& definition : result.definitions_) { + if (definition.index_type == parts[0] && definition.column_name == parts[1]) { + definition.options[parts[2]] = value; + found = true; + break; + } + } + if (!found) { + return Status::Invalid( + fmt::format("Wrong file index option '{}': column '{}' is not declared in " + "'file-index.{}.columns'", + key, parts[1], parts[0])); + } + } + return result; +} + +} // namespace paimon diff --git a/src/paimon/core/io/file_index_options.h b/src/paimon/core/io/file_index_options.h new file mode 100644 index 000000000..ae51f05d4 --- /dev/null +++ b/src/paimon/core/io/file_index_options.h @@ -0,0 +1,61 @@ +/* + * 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 +#include +#include +#include + +#include "paimon/result.h" + +namespace paimon { + +class CoreOptions; + +struct FileIndexDefinition { + std::string column_name; + std::string index_type; + std::map options; +}; + +/// Parsed write-side file index configuration. +class FileIndexOptions { + public: + static Result FromCoreOptions(const CoreOptions& options); + + const std::vector& Definitions() const { + return definitions_; + } + + int64_t InManifestThreshold() const { + return in_manifest_threshold_; + } + + bool Empty() const { + return definitions_.empty(); + } + + private: + std::vector definitions_; + int64_t in_manifest_threshold_ = 0; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/key_value_data_file_writer.cpp b/src/paimon/core/io/key_value_data_file_writer.cpp index 9393c7c3c..56a8bcd4b 100644 --- a/src/paimon/core/io/key_value_data_file_writer.cpp +++ b/src/paimon/core/io/key_value_data_file_writer.cpp @@ -25,7 +25,6 @@ #include #include -#include "arrow/type.h" #include "fmt/format.h" #include "paimon/common/data/binary_array.h" #include "paimon/common/data/binary_array_writer.h" @@ -53,7 +52,7 @@ KeyValueDataFileWriter::KeyValueDataFileWriter( const std::shared_ptr& stats_extractor, const std::shared_ptr& write_schema, bool is_external_path, const std::shared_ptr& pool) - : SingleFileWriter(compression, converter), + : DataFileWriterBase(compression, std::move(converter)), pool_(pool), schema_id_(schema_id), level_(level), @@ -64,10 +63,6 @@ KeyValueDataFileWriter::KeyValueDataFileWriter( is_external_path_(is_external_path), disable_stats_(stats_extractor == nullptr) {} -void KeyValueDataFileWriter::SetMetadataFinalizer(MetadataFinalizer finalizer) { - metadata_finalizer_ = std::move(finalizer); -} - Status KeyValueDataFileWriter::Write(KeyValueBatch batch) { // update min and max key if (!min_key_) { @@ -80,19 +75,8 @@ Status KeyValueDataFileWriter::Write(KeyValueBatch batch) { // update delete row count delete_row_count_ += batch.delete_row_count; - PAIMON_RETURN_NOT_OK(SingleFileWriter::Write(std::move(batch))); - return Status::OK(); -} - -Status KeyValueDataFileWriter::BeforeFinish() { - if (metadata_finalizer_) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr updated_schema, - metadata_finalizer_()); - if (updated_schema) { - PAIMON_RETURN_NOT_OK(UpdateSchema(updated_schema)); - } - } - return Status::OK(); + ::ArrowArray* logical_batch = batch.batch.get(); + return WriteRecord(std::move(batch), logical_batch); } Result> KeyValueDataFileWriter::GetResult() { @@ -120,14 +104,14 @@ Result> KeyValueDataFileWriter::GetResult() { final_path = external_path.ToString(); } PAIMON_ASSIGN_OR_RAISE(int64_t local_micro, DateTimeUtils::GetCurrentLocalTimeUs()); + const FileIndexWriteResult& file_index = GetFileIndexWriteResult(); return std::make_shared( PathUtil::GetName(path_), output_bytes_, RecordCount(), min_key, max_key, key_stats, value_stats, min_sequence_number_, max_sequence_number_, schema_id_, level_, - /*extra_files=*/std::vector>(), + file_index.extra_files, Timestamp(/*millisecond=*/local_micro / 1000, /*nano_of_millisecond=*/0), delete_row_count_, - /*embedded_index=*/nullptr, file_source_, - /*value_stats_cols=*/std::nullopt, final_path, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + file_index.embedded_index, file_source_, /*value_stats_cols=*/std::nullopt, final_path, + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); } Status KeyValueDataFileWriter::GenerateMinMaxKey(BinaryRow* min_key, BinaryRow* max_key) const { diff --git a/src/paimon/core/io/key_value_data_file_writer.h b/src/paimon/core/io/key_value_data_file_writer.h index e1e3fd92c..eb7a2efca 100644 --- a/src/paimon/core/io/key_value_data_file_writer.h +++ b/src/paimon/core/io/key_value_data_file_writer.h @@ -17,6 +17,7 @@ */ #pragma once + #include #include #include @@ -25,7 +26,7 @@ #include #include "paimon/core/io/data_file_meta.h" -#include "paimon/core/io/single_file_writer.h" +#include "paimon/core/io/data_file_writer_base.h" #include "paimon/core/key_value.h" #include "paimon/core/manifest/file_source.h" #include "paimon/result.h" @@ -44,14 +45,8 @@ class InternalRow; class MemoryPool; class SimpleStats; -class KeyValueDataFileWriter - : public SingleFileWriter> { +class KeyValueDataFileWriter : public DataFileWriterBase { public: - /// Callback invoked during BeforeFinish() to finalize file metadata. - /// Produces an updated schema with per-field metadata (e.g. shredding metadata) - /// and may perform other finalization work (e.g. reporting stats to cross-file context). - using MetadataFinalizer = std::function>()>; - KeyValueDataFileWriter(const std::string& compression, std::function converter, int64_t schema_id, int32_t level, FileSource file_source, @@ -60,17 +55,10 @@ class KeyValueDataFileWriter const std::shared_ptr& write_schema, bool is_external_path, const std::shared_ptr& pool); - /// Sets the metadata finalizer. Called during BeforeFinish() to produce an updated - /// schema and perform finalization callbacks. Must be set before Close(). - void SetMetadataFinalizer(MetadataFinalizer finalizer); - Status Write(KeyValueBatch batch) override; Result> GetResult() override; - protected: - Status BeforeFinish() override; - private: Result>> GetFieldStats(); @@ -96,7 +84,6 @@ class KeyValueDataFileWriter int64_t max_sequence_number_ = std::numeric_limits::min(); std::shared_ptr min_key_; std::shared_ptr max_key_; - MetadataFinalizer metadata_finalizer_; }; } // namespace paimon diff --git a/src/paimon/core/io/key_value_data_file_writer_factory.cpp b/src/paimon/core/io/key_value_data_file_writer_factory.cpp index 07d50b980..8f3885591 100644 --- a/src/paimon/core/io/key_value_data_file_writer_factory.cpp +++ b/src/paimon/core/io/key_value_data_file_writer_factory.cpp @@ -24,6 +24,7 @@ #include "arrow/c/helpers.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/key_value_data_file_writer.h" #include "paimon/format/file_format.h" @@ -60,6 +61,11 @@ KeyValueDataFileWriterFactory::CreateWriter() const { options_.GetWriteFileCompression(level_), std::move(converter), schema_id_, level_, file_source_, primary_keys_, resources.stats_extractor, write_schema_, path_factory_->IsExternalPath(), pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_index_writer, + CreateFileIndexWriter(write_schema_, path_factory_)); + if (file_index_writer) { + writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); return std::unique_ptr>>( diff --git a/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp b/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp index 6e4843bbb..0e4e82199 100644 --- a/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp +++ b/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp @@ -23,6 +23,7 @@ #include "arrow/c/helpers.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/data_file_writer.h" #include "paimon/core/io/infer_shredding_file_writer.h" @@ -89,6 +90,11 @@ ShreddingAppendDataFileWriterFactory::CreateShreddedWriter( options_.GetFileCompression(), std::move(batch_converter), schema_id_, seq_num_counter, file_source_, resources.stats_extractor, path_factory_->IsExternalPath(), write_cols_, pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_index_writer, + CreateFileIndexWriter(write_schema_, path_factory_)); + if (file_index_writer) { + writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); ShreddingWritePlanFactory::MetadataFinalizer finalizer = diff --git a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp index 8ac583ee0..30d4c9fce 100644 --- a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp +++ b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp @@ -23,6 +23,7 @@ #include "arrow/c/helpers.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/infer_shredding_file_writer.h" #include "paimon/core/io/key_value_data_file_writer.h" @@ -88,6 +89,11 @@ ShreddingKeyValueDataFileWriterFactory::CreateShreddedWriter( options_.GetWriteFileCompression(level_), std::move(batch_converter), schema_id_, level_, file_source_, primary_keys_, resources.stats_extractor, file_schema, path_factory_->IsExternalPath(), pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_index_writer, + CreateFileIndexWriter(write_schema_, path_factory_)); + if (file_index_writer) { + writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); ShreddingWritePlanFactory::MetadataFinalizer finalizer = diff --git a/src/paimon/core/io/single_file_writer.h b/src/paimon/core/io/single_file_writer.h index 99507b579..6db3a699c 100644 --- a/src/paimon/core/io/single_file_writer.h +++ b/src/paimon/core/io/single_file_writer.h @@ -25,6 +25,7 @@ #include #include #include +#include #include "arrow/c/abi.h" #include "arrow/c/bridge.h" @@ -64,21 +65,27 @@ class SingleFileWriter : public FileWriter { class AbortExecutor { public: AbortExecutor(const std::shared_ptr& fs, const std::string& path) - : fs_(fs), path_(path), logger_(Logger::GetLogger("AbortExecutor")) {} + : paths_({{fs, path}}), logger_(Logger::GetLogger("AbortExecutor")) {} + + void Add(const std::shared_ptr& fs, const std::string& path) { + paths_.emplace_back(fs, path); + } void Abort() { - if (fs_) { - auto status = fs_->Delete(path_); + for (const auto& [fs, path] : paths_) { + if (!fs) { + continue; + } + auto status = fs->Delete(path); if (!status.ok()) { - PAIMON_LOG_WARN(logger_, "Exception occurs when deleting %s: %s", path_.c_str(), + PAIMON_LOG_WARN(logger_, "Exception occurs when deleting %s: %s", path.c_str(), status.ToString().c_str()); } } } private: - std::shared_ptr fs_; - std::string path_; + std::vector, std::string>> paths_; std::shared_ptr logger_; }; @@ -132,6 +139,11 @@ class SingleFileWriter : public FileWriter { return Status::OK(); } + /// Hook called after the data file is closed and before its completion callback is published. + virtual Status BeforeCompletion() { + return Status::OK(); + } + /// Serializes schema and forwards it as file metadata to FormatWriter. Status UpdateSchema(const std::shared_ptr& schema); @@ -239,6 +251,7 @@ Status SingleFileWriter::Close() { // guard still removes the file on a callback error, while a repeated Close() does not publish // the same file again. closed_ = true; + PAIMON_RETURN_NOT_OK(BeforeCompletion()); if (completion_callback_) { PAIMON_RETURN_NOT_OK(completion_callback_()); } From 2fdd89593230298bea02605acd5fe5314df3af65 Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Mon, 17 Aug 2026 19:18:10 +0800 Subject: [PATCH 2/4] test(file-index): add end-to-end write coverage --- test/inte/write_and_read_inte_test.cpp | 83 ++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp index 82eb5257b..5916fc97c 100644 --- a/test/inte/write_and_read_inte_test.cpp +++ b/test/inte/write_and_read_inte_test.cpp @@ -486,6 +486,89 @@ TEST_P(WriteAndReadInteTest, TestAppendVectorWithPredicate) { ASSERT_TRUE(expected->Equals(actual)) << actual->ToString(); } +// TODO(jinli.zjw): move to a single file for a file index inte test +TEST_P(WriteAndReadInteTest, TestAppendWithExternalBitmapAndRangeBitmapIndexes) { + arrow::FieldVector fields = {arrow::field("name", arrow::utf8()), + arrow::field("score", arrow::int32())}; + auto [file_format, file_system] = GetParam(); + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, + {Options::FILE_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1MB"}, + {Options::BUCKET, "-1"}, + {Options::FILE_SYSTEM, file_system}, + {"file-index.bitmap.columns", "name"}, + {"file-index.range-bitmap.columns", "score"}, + {"file-index.range-bitmap.score.chunk-size", "1KB"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1B"}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + + auto schema = arrow::schema(fields); + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(test_dir_, schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([ + ["Alice", 10], + ["Bob", 20], + ["Alice", 30], + ["Lucy", 40] + ])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN(auto data_files, CurrentDataFiles(options)); + ASSERT_EQ(1, data_files.size()); + const auto& [bucket_path, data_file] = data_files[0]; + ASSERT_FALSE(data_file->embedded_index); + ASSERT_EQ(1, data_file->extra_files.size()); + ASSERT_TRUE(data_file->extra_files[0]); + ASSERT_EQ(data_file->file_name + ".index", data_file->extra_files[0].value()); + std::string index_path = PathUtil::JoinPath(bucket_path, data_file->extra_files[0].value()); + ASSERT_OK_AND_ASSIGN(bool index_exists, dir_->GetFileSystem()->Exists(index_path)); + ASSERT_TRUE(index_exists); + + std::string indexed_name = "Alice"; + auto name_predicate = PredicateBuilder::Equal( + /*field_index=*/0, /*field_name=*/"name", FieldType::STRING, + Literal(FieldType::STRING, indexed_name.data(), indexed_name.size())); + auto score_predicate = PredicateBuilder::GreaterThan( + /*field_index=*/1, /*field_name=*/"score", FieldType::INT, Literal(20)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({name_predicate, score_predicate})); + + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ScanContextBuilder scan_context_builder(table_path); + scan_context_builder.SetOptions(options) + .AddOption(Options::SCAN_MODE, StartupMode::LatestFull().ToString()) + .SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto scan_context, scan_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(auto plan, table_scan->CreatePlan()); + ASSERT_EQ(1, plan->Splits().size()); + + // Keep precise post-read filtering disabled. The exact result therefore verifies that the + // bitmap and range-bitmap indexes produced by the write path are consumed by the read path. + ReadContextBuilder read_context_builder(table_path); + read_context_builder.SetOptions(options).SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + + arrow::FieldVector fields_with_row_kind = fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + auto expected_result = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(fields_with_row_kind), R"([[0, "Alice", 30]])"); + ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); + auto expected = std::make_shared(expected_result.ValueOrDie()); + ASSERT_TRUE(expected->Equals(actual)) << actual->ToString(); +} + TEST_P(WriteAndReadInteTest, TestPKSimple) { arrow::FieldVector fields = { arrow::field("pk", arrow::utf8()), From 2ace9ed1cf17659b5800ec3f7b56627bc6286d75 Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Tue, 18 Aug 2026 17:15:31 +0800 Subject: [PATCH 3/4] refactor(file-index): refine index serialization --- include/paimon/file_index/file_index_format.h | 7 +- .../common/file_index/file_index_format.cpp | 110 +++++++++++------- .../file_index/file_index_format_test.cpp | 17 +-- .../common/io/byte_array_output_stream.cpp | 27 ++--- .../common/io/byte_array_output_stream.h | 12 +- .../io/byte_array_output_stream_test.cpp | 29 ++--- .../io/data_input_output_stream_test.cpp | 8 +- .../io/memory_segment_output_stream.cpp | 16 +-- src/paimon/core/io/data_file_index_writer.cpp | 10 +- .../core/io/data_file_index_writer_test.cpp | 3 +- src/paimon/core/io/data_file_writer_base.h | 6 +- 11 files changed, 129 insertions(+), 116 deletions(-) diff --git a/include/paimon/file_index/file_index_format.h b/include/paimon/file_index/file_index_format.h index 38459845a..e2d581e52 100644 --- a/include/paimon/file_index/file_index_format.h +++ b/include/paimon/file_index/file_index_format.h @@ -92,10 +92,14 @@ class PAIMON_EXPORT FileIndexFormat { public: class Reader; class Writer; + + /// Serialized file indexes grouped as column name -> index type -> index bytes. A null bytes + /// pointer represents an empty index for that column and index type. + /// For example, indexes["col1"]["bsi"] = ; using ColumnIndexes = std::map>>; + /// Creates a `Reader` to parse a index file (may contain multiple indexes) from the given input /// stream. - /// /// @param input_stream Input stream containing serialized index data. /// @param pool Memory pool for temporary allocations during reading. /// @return A unique pointer to a `Reader` on success, or an error if the stream is invalid @@ -130,7 +134,6 @@ class FileIndexFormat::Reader { public: virtual ~Reader() = default; /// Reads index data for a specific column from the index file. - /// /// @param column_name Name of the column to retrieve index data for. /// @param arrow_schema Arrow schema that must contain a field corresponding to `column_name`. /// @return A vector of shared pointers to FileIndexReader objects, each corresponding to a diff --git a/src/paimon/common/file_index/file_index_format.cpp b/src/paimon/common/file_index/file_index_format.cpp index 23bf3715c..fab5c7a7b 100644 --- a/src/paimon/common/file_index/file_index_format.cpp +++ b/src/paimon/common/file_index/file_index_format.cpp @@ -52,81 +52,107 @@ class FileIndexFormatWriterImpl : public FileIndexFormat::Writer { if (written_) { return Status::Invalid("File index column indexes have already been written"); } - int64_t header_length = sizeof(int64_t) + sizeof(int32_t) * 3 + sizeof(int32_t); + + PAIMON_RETURN_NOT_OK(WriteHead(indexes)); + // Write body. + DataOutputStream data_output(output_stream_); + for (const auto& [column_name, column_indexes] : indexes) { + for (const auto& [index_type, bytes] : column_indexes) { + if (bytes) { + PAIMON_RETURN_NOT_OK(data_output.WriteBytes(bytes)); + } + } + } + written_ = true; + return Status::OK(); + } + + Status Close() override { + if (closed_) { + return Status::OK(); + } + closed_ = true; + PAIMON_RETURN_NOT_OK(output_stream_->Flush()); + return output_stream_->Close(); + } + + private: + static constexpr int32_t kRedundantLength = 0; + + static Result CalculateHeadLength(const FileIndexFormat::ColumnIndexes& indexes) { + // magic(8), version(4), header length(4), and column count(4). + int64_t head_length = 8 + 4 + 4 + 4; int64_t body_length = 0; + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(indexes.size(), "file index column count")); for (const auto& [column_name, column_indexes] : indexes) { PAIMON_RETURN_NOT_OK(ValidateValueInRange(column_name.size(), "file index column name length")); - header_length += - sizeof(uint16_t) + static_cast(column_name.size()) + sizeof(int32_t); + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(column_indexes.size(), "column index count")); + // column name(2 + N) + index count(4) + head_length += 2 + static_cast(column_name.size()) + 4; for (const auto& [index_type, bytes] : column_indexes) { PAIMON_RETURN_NOT_OK(ValidateValueInRange(index_type.size(), "file index type name length")); - header_length += sizeof(uint16_t) + static_cast(index_type.size()) + - sizeof(int32_t) * 2; + // index type(2 + N) + body offset(4) + body length(4) + head_length += 2 + static_cast(index_type.size()) + 4 + 4; if (bytes) { - PAIMON_RETURN_NOT_OK(AddChecked(bytes->size(), &body_length, "index body")); + PAIMON_RETURN_NOT_OK(AddChecked(bytes->size(), "index body", &body_length)); } } } + + head_length += 4; // The trailing redundant-length field(4). PAIMON_RETURN_NOT_OK( - ValidateValueInRange(header_length, "file index header length")); - PAIMON_RETURN_NOT_OK( - ValidateValueInRange(indexes.size(), "file index column count")); - PAIMON_RETURN_NOT_OK(AddChecked(header_length, &body_length, "file index size")); + ValidateValueInRange(head_length, "file index header length")); + int64_t container_length = head_length + body_length; + PAIMON_RETURN_NOT_OK(ValidateValueInRange(container_length, "file index size")); + return static_cast(head_length); + } + Status WriteHead(const FileIndexFormat::ColumnIndexes& indexes) { + PAIMON_ASSIGN_OR_RAISE(int32_t head_length, CalculateHeadLength(indexes)); DataOutputStream data_output(output_stream_); + // Write magic. PAIMON_RETURN_NOT_OK(data_output.WriteValue(FileIndexFormat::MAGIC)); + // Write version. PAIMON_RETURN_NOT_OK(data_output.WriteValue(FileIndexFormat::V_1)); - PAIMON_RETURN_NOT_OK(data_output.WriteValue(static_cast(header_length))); + // Write head length. + PAIMON_RETURN_NOT_OK(data_output.WriteValue(head_length)); + // Write column count. PAIMON_RETURN_NOT_OK(data_output.WriteValue(static_cast(indexes.size()))); - int64_t body_offset = header_length; + int64_t body_offset = head_length; for (const auto& [column_name, column_indexes] : indexes) { + // Write column name. PAIMON_RETURN_NOT_OK(data_output.WriteString(column_name)); - PAIMON_RETURN_NOT_OK( - ValidateValueInRange(column_indexes.size(), "column index count")); + // Write index count for the column. PAIMON_RETURN_NOT_OK( data_output.WriteValue(static_cast(column_indexes.size()))); for (const auto& [index_type, bytes] : column_indexes) { + // Write index type. PAIMON_RETURN_NOT_OK(data_output.WriteString(index_type)); - if (bytes == nullptr) { + // Write body offset and length. + if (bytes) { + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(static_cast(body_offset))); + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(static_cast(bytes->size()))); + body_offset += static_cast(bytes->size()); + } else { PAIMON_RETURN_NOT_OK( data_output.WriteValue(FileIndexFormat::EMPTY_INDEX_FLAG)); PAIMON_RETURN_NOT_OK(data_output.WriteValue(0)); - continue; - } - PAIMON_RETURN_NOT_OK( - data_output.WriteValue(static_cast(body_offset))); - PAIMON_RETURN_NOT_OK( - data_output.WriteValue(static_cast(bytes->size()))); - body_offset += static_cast(bytes->size()); - } - } - PAIMON_RETURN_NOT_OK(data_output.WriteValue(0)); - for (const auto& [column_name, column_indexes] : indexes) { - for (const auto& [index_type, bytes] : column_indexes) { - if (bytes) { - PAIMON_RETURN_NOT_OK(data_output.WriteBytes(bytes)); } } } - written_ = true; - return Status::OK(); - } - - Status Close() override { - if (closed_) { - return Status::OK(); - } - closed_ = true; - PAIMON_RETURN_NOT_OK(output_stream_->Flush()); - return output_stream_->Close(); + // Write redundant length for future format extensions. + return data_output.WriteValue(kRedundantLength); } - private: template - static Status AddChecked(T value, int64_t* total, const char* name) { + static Status AddChecked(T value, const char* name, int64_t* total) { PAIMON_RETURN_NOT_OK(ValidateValueInRange(value, name)); *total += static_cast(value); return ValidateValueInRange(*total, name); diff --git a/src/paimon/common/file_index/file_index_format_test.cpp b/src/paimon/common/file_index/file_index_format_test.cpp index 148c71b3d..40989f7e9 100644 --- a/src/paimon/common/file_index/file_index_format_test.cpp +++ b/src/paimon/common/file_index/file_index_format_test.cpp @@ -58,32 +58,25 @@ class FileIndexFormatTest : public ::testing::Test { std::shared_ptr pool_; }; -TEST_F(FileIndexFormatTest, TestWriteEmptyIndexGoldenBytes) { +TEST_F(FileIndexFormatTest, TestWriteAndReadEmptyIndexGoldenBytes) { // the expected bytes are generated from Java Paimon std::vector expected = {0, 5, 78, 78, -48, 26, 53, -82, 0, 0, 0, 1, 0, 0, 0, 47, 0, 0, 0, 1, 0, 2, 99, 49, 0, 0, 0, 1, 0, 5, 101, 109, 112, 116, 121, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0}; FileIndexFormat::ColumnIndexes indexes; indexes["c1"]["empty"] = nullptr; - auto output = std::make_shared( + auto segment_output = std::make_unique( MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_); + auto output = std::make_shared(std::move(segment_output)); ASSERT_OK_AND_ASSIGN(auto writer, FileIndexFormat::CreateWriter(output, pool_)); ASSERT_OK(writer->WriteColumnIndexes(indexes)); ASSERT_OK(writer->Close()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, output->Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, output->Finish(pool_.get())); ASSERT_EQ(expected, std::vector(actual->data(), actual->data() + actual->size())); -} - -TEST_F(FileIndexFormatTest, TestCreateEmptyFileIndexReader) { auto schema = arrow::schema({arrow::field("c1", arrow::utf8())}); - std::vector index_file_bytes = {0, 5, 78, 78, -48, 26, 53, -82, 0, 0, 0, 1, - 0, 0, 0, 47, 0, 0, 0, 1, 0, 2, 99, 49, - 0, 0, 0, 1, 0, 5, 101, 109, 112, 116, 121, -1, - -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0}; - auto input_stream = - std::make_shared(index_file_bytes.data(), index_file_bytes.size()); + auto input_stream = std::make_shared(actual->data(), actual->size()); ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input_stream, pool_)); ASSERT_OK_AND_ASSIGN(auto index_file_readers, reader->ReadColumnIndex("c1", CreateArrowSchema(schema).get())); diff --git a/src/paimon/common/io/byte_array_output_stream.cpp b/src/paimon/common/io/byte_array_output_stream.cpp index 94da8deae..bc6d1f1ab 100644 --- a/src/paimon/common/io/byte_array_output_stream.cpp +++ b/src/paimon/common/io/byte_array_output_stream.cpp @@ -19,7 +19,9 @@ #include "paimon/common/io/byte_array_output_stream.h" #include +#include #include +#include #include #include "paimon/common/memory/memory_segment_utils.h" @@ -29,9 +31,10 @@ namespace paimon { -ByteArrayOutputStream::ByteArrayOutputStream(int32_t initial_capacity, - const std::shared_ptr& pool) - : pool_(pool), output_(initial_capacity, pool_) {} +ByteArrayOutputStream::ByteArrayOutputStream(std::unique_ptr&& output) + : output_(std::move(output)) { + assert(output_); +} Result ByteArrayOutputStream::Write(const char* buffer, int64_t size) { if (closed_) { @@ -45,11 +48,10 @@ Result ByteArrayOutputStream::Write(const char* buffer, int64_t size) { while (remaining > 0) { uint32_t to_write = static_cast(std::min( remaining, static_cast(std::numeric_limits::max()))); - output_.Write(buffer, to_write); + output_->Write(buffer, to_write); buffer += to_write; remaining -= to_write; } - position_ += size; return size; } @@ -58,21 +60,20 @@ Status ByteArrayOutputStream::Close() { return Status::OK(); } -Result> ByteArrayOutputStream::Finish() { +Result> ByteArrayOutputStream::Finish(MemoryPool* pool) { + assert(pool); PAIMON_RETURN_NOT_OK(Close()); if (result_) { return result_; } // TODO(jinli.zjw): Support int64_t lengths in MemorySegmentUtils::CopyToBytes and remove this // limit. - if (position_ > std::numeric_limits::max()) { - return Status::Invalid("Byte array output stream size exceeds INT32_MAX"); - } - const std::vector& segments = output_.Segments(); - result_ = std::shared_ptr(new Bytes(static_cast(position_), pool_.get()), - [pool = pool_](Bytes* bytes) { delete bytes; }); + const int64_t size = output_->CurrentSize(); + PAIMON_RETURN_NOT_OK(ValidateValueInRange(size, "byte array output stream size")); + const std::vector& segments = output_->Segments(); + result_ = std::make_shared(static_cast(size), pool); MemorySegmentUtils::CopyToBytes(segments, /*offset=*/0, result_.get(), - /*bytes_offset=*/0, static_cast(position_)); + /*bytes_offset=*/0, static_cast(size)); return result_; } diff --git a/src/paimon/common/io/byte_array_output_stream.h b/src/paimon/common/io/byte_array_output_stream.h index bceec378d..9b87ca429 100644 --- a/src/paimon/common/io/byte_array_output_stream.h +++ b/src/paimon/common/io/byte_array_output_stream.h @@ -35,7 +35,8 @@ class MemoryPool; /// An in-memory output stream backed by segments allocated from a Paimon MemoryPool. class ByteArrayOutputStream : public OutputStream { public: - ByteArrayOutputStream(int32_t initial_capacity, const std::shared_ptr& pool); + /// Takes ownership of an initialized segmented output stream. + explicit ByteArrayOutputStream(std::unique_ptr&& output); ~ByteArrayOutputStream() override = default; @@ -46,7 +47,7 @@ class ByteArrayOutputStream : public OutputStream { } Result GetPos() const override { - return position_; + return output_->CurrentSize(); } Result GetUri() const override { @@ -56,13 +57,12 @@ class ByteArrayOutputStream : public OutputStream { Status Close() override; /// Closes the stream and returns its contents as an exactly-sized contiguous byte array. - Result> Finish(); + /// @note The caller must keep `pool` alive until the returned bytes are destroyed. + Result> Finish(MemoryPool* pool); private: - std::shared_ptr pool_; - MemorySegmentOutputStream output_; + std::unique_ptr output_; std::shared_ptr result_; - int64_t position_ = 0; bool closed_ = false; }; diff --git a/src/paimon/common/io/byte_array_output_stream_test.cpp b/src/paimon/common/io/byte_array_output_stream_test.cpp index 23ce567aa..bd185095d 100644 --- a/src/paimon/common/io/byte_array_output_stream_test.cpp +++ b/src/paimon/common/io/byte_array_output_stream_test.cpp @@ -20,6 +20,7 @@ #include #include +#include #include "gtest/gtest.h" #include "paimon/memory/bytes.h" @@ -30,8 +31,9 @@ namespace paimon::test { TEST(ByteArrayOutputStreamTest, TestWriteAndFinish) { std::shared_ptr pool = GetMemoryPool(); + auto output = std::make_unique(/*segment_size=*/2, pool); std::shared_ptr stream = - std::make_shared(/*initial_capacity=*/2, pool); + std::make_shared(std::move(output)); ASSERT_GT(pool->CurrentUsage(), 0); ASSERT_OK_AND_ASSIGN(int64_t first_write, stream->Write("ab", 2)); ASSERT_EQ(2, first_write); @@ -41,43 +43,44 @@ TEST(ByteArrayOutputStreamTest, TestWriteAndFinish) { ASSERT_EQ(6, position); ASSERT_EQ(pool->CurrentUsage(), pool->MaxMemoryUsage()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); ASSERT_EQ("abcdef", std::string(result->data(), result->size())); ASSERT_NOK_WITH_MSG(stream->Write("x", 1), "closed"); - ASSERT_OK_AND_ASSIGN(std::shared_ptr repeated, stream->Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr repeated, stream->Finish(pool.get())); ASSERT_EQ(result, repeated); stream.reset(); ASSERT_EQ(6, pool->CurrentUsage()); } TEST(ByteArrayOutputStreamTest, TestWriteValidation) { - std::shared_ptr stream = std::make_shared( - /*initial_capacity=*/8, GetDefaultPool()); + std::shared_ptr pool = GetDefaultPool(); + auto output = std::make_unique(/*segment_size=*/8, pool); + std::shared_ptr stream = + std::make_shared(std::move(output)); ASSERT_NOK(stream->Write(nullptr, 1)); ASSERT_NOK(stream->Write("", -1)); ASSERT_OK_AND_ASSIGN(int64_t written, stream->Write(nullptr, 0)); ASSERT_EQ(0, written); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); ASSERT_EQ(0, result->size()); } -TEST(ByteArrayOutputStreamTest, TestResultKeepsMemoryPoolAlive) { +TEST(ByteArrayOutputStreamTest, TestCallerKeepsMemoryPoolAlive) { std::shared_ptr pool = GetMemoryPool(); - std::weak_ptr weak_pool = pool; + auto output = std::make_unique(/*segment_size=*/8, pool); std::shared_ptr stream = - std::make_shared(/*initial_capacity=*/8, pool); + std::make_shared(std::move(output)); ASSERT_OK_AND_ASSIGN(int64_t written, stream->Write("data", 4)); ASSERT_EQ(4, written); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); stream.reset(); - pool.reset(); - ASSERT_FALSE(weak_pool.expired()); + ASSERT_GT(pool->CurrentUsage(), 0); ASSERT_EQ("data", std::string(result->data(), result->size())); result.reset(); - ASSERT_TRUE(weak_pool.expired()); + ASSERT_EQ(0, pool->CurrentUsage()); } } // namespace paimon::test diff --git a/src/paimon/common/io/data_input_output_stream_test.cpp b/src/paimon/common/io/data_input_output_stream_test.cpp index 4e6063706..0a5dd5748 100644 --- a/src/paimon/common/io/data_input_output_stream_test.cpp +++ b/src/paimon/common/io/data_input_output_stream_test.cpp @@ -79,12 +79,8 @@ class DataInputOutputStreamTest : public ::testing::Test, (void)data_output_stream->WriteValue(static_cast(9223372036854775805)); // 8 bytes (void)data_output_stream->WriteValue(true); // 1 byte std::string str1 = "This is a very very very long sentence."; - if constexpr (std::is_same_v) { - (void)data_output_stream->WriteString(str1); // 39 bytes + 2 bytes len - } else { - (void)data_output_stream->WriteString(str1); // 39 bytes + 2 bytes len - } - std::string str2 = "我是一个粉刷匠～"; // 24 bytes + (void)data_output_stream->WriteString(str1); // 39 bytes + 2 bytes len + std::string str2 = "我是一个粉刷匠～"; // 24 bytes auto bytes = std::make_shared(str2, pool_.get()); (void)data_output_stream->WriteBytes(bytes); } diff --git a/src/paimon/common/io/memory_segment_output_stream.cpp b/src/paimon/common/io/memory_segment_output_stream.cpp index cbcfabf9c..2d0d274a7 100644 --- a/src/paimon/common/io/memory_segment_output_stream.cpp +++ b/src/paimon/common/io/memory_segment_output_stream.cpp @@ -54,20 +54,8 @@ void MemorySegmentOutputStream::WriteString(const std::string& str) { } void MemorySegmentOutputStream::Write(const char* data, uint32_t size) { - uint32_t remaining_size = size; - while (remaining_size > 0) { - int32_t remaining_in_segment = segment_size_ - position_in_segment_; - if (remaining_in_segment == 0) { - Advance(); - remaining_in_segment = segment_size_; - } - int32_t to_write = static_cast( - std::min(remaining_size, static_cast(remaining_in_segment))); - std::memcpy(current_segment_.MutableData() + position_in_segment_, data, to_write); - data += to_write; - remaining_size -= to_write; - position_in_segment_ += to_write; - } + MemorySegment segment = MemorySegment::WrapView(data, size); + Write(segment, 0, segment.Size()); } void MemorySegmentOutputStream::Write(const MemorySegment& segment, int32_t offset, int32_t len) { diff --git a/src/paimon/core/io/data_file_index_writer.cpp b/src/paimon/core/io/data_file_index_writer.cpp index 844a30e4a..2d21a02b1 100644 --- a/src/paimon/core/io/data_file_index_writer.cpp +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -92,7 +92,7 @@ DataFileIndexWriter::DataFileIndexWriter(std::vector&& writers pool_(pool) {} Status DataFileIndexWriter::AddBatch(const std::shared_ptr& logical_batch) { - for (IndexWriterEntry& entry : writers_) { + for (const IndexWriterEntry& entry : writers_) { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::shared_ptr projected, arrow::StructArray::Make({logical_batch->field(entry.field_index)}, {entry.field})); @@ -107,20 +107,22 @@ Status DataFileIndexWriter::AddBatch(const std::shared_ptr& Result> DataFileIndexWriter::SerializeContainer() { FileIndexFormat::ColumnIndexes column_indexes; - for (IndexWriterEntry& entry : writers_) { + for (const IndexWriterEntry& entry : writers_) { PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR serialized, entry.writer->SerializedBytes()); column_indexes[entry.column_name][entry.index_type] = std::shared_ptr(std::move(serialized)); } - std::shared_ptr output = std::make_shared( + auto segment_output = std::make_unique( MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_); + std::shared_ptr output = + std::make_shared(std::move(segment_output)); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_writer, FileIndexFormat::CreateWriter(output, pool_)); PAIMON_RETURN_NOT_OK(format_writer->WriteColumnIndexes(column_indexes)); PAIMON_RETURN_NOT_OK(format_writer->Close()); - return output->Finish(); + return output->Finish(pool_.get()); } Result DataFileIndexWriter::Finish(const std::string& data_file_path) { diff --git a/src/paimon/core/io/data_file_index_writer_test.cpp b/src/paimon/core/io/data_file_index_writer_test.cpp index 5e881b4e1..fc9b02e1c 100644 --- a/src/paimon/core/io/data_file_index_writer_test.cpp +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -28,6 +28,7 @@ #include "arrow/type.h" #include "gtest/gtest.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/file_index_options.h" @@ -70,7 +71,7 @@ class DataFileIndexWriterTest : public ::testing::Test { std::shared_ptr array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema_->fields()), json) .ValueOrDie(); - return std::dynamic_pointer_cast(array); + return checked_pointer_cast(array); } Result> CreateReader( diff --git a/src/paimon/core/io/data_file_writer_base.h b/src/paimon/core/io/data_file_writer_base.h index 64b556857..260402081 100644 --- a/src/paimon/core/io/data_file_writer_base.h +++ b/src/paimon/core/io/data_file_writer_base.h @@ -25,7 +25,7 @@ #include "arrow/c/bridge.h" #include "arrow/type.h" -#include "arrow/util/checked_cast.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/single_file_writer.h" @@ -57,7 +57,7 @@ class DataFileWriterBase : public SingleFileWriter file_index_writer, + void SetFileIndexWriter(std::unique_ptr&& file_index_writer, const std::shared_ptr& logical_schema) { file_index_writer_ = std::move(file_index_writer); logical_type_ = arrow::struct_(logical_schema->fields()); @@ -119,7 +119,7 @@ class DataFileWriterBase : public SingleFileWriter logical_array, arrow::ImportArray(batch, logical_type_)); std::shared_ptr logical_batch = - arrow::internal::checked_pointer_cast(logical_array); + checked_pointer_cast(logical_array); PAIMON_RETURN_NOT_OK(file_index_writer_->AddBatch(logical_batch)); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*logical_batch, batch)); return Status::OK(); From e3fe5c73e92853d4780a0ab72785dc3e5477085f Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Wed, 19 Aug 2026 16:39:36 +0800 Subject: [PATCH 4/4] fix(file-index): refine writer validation and lifecycle --- include/paimon/file_index/file_index_format.h | 5 + src/paimon/CMakeLists.txt | 1 + src/paimon/core/io/data_file_index_writer.cpp | 23 ++++- src/paimon/core/io/data_file_index_writer.h | 7 ++ .../core/io/data_file_index_writer_test.cpp | 97 ++++++++++++++++++- src/paimon/core/io/data_file_writer.cpp | 2 +- src/paimon/core/io/data_file_writer_base.h | 23 +++-- src/paimon/core/io/file_index_options.cpp | 21 ++-- src/paimon/core/io/file_index_options.h | 2 + .../core/io/file_index_options_test.cpp | 58 +++++++++++ .../core/io/key_value_data_file_writer.cpp | 3 +- 11 files changed, 218 insertions(+), 24 deletions(-) create mode 100644 src/paimon/core/io/file_index_options_test.cpp diff --git a/include/paimon/file_index/file_index_format.h b/include/paimon/file_index/file_index_format.h index e2d581e52..3993b6247 100644 --- a/include/paimon/file_index/file_index_format.h +++ b/include/paimon/file_index/file_index_format.h @@ -108,6 +108,10 @@ class PAIMON_EXPORT FileIndexFormat { const std::shared_ptr& input_stream, const std::shared_ptr& pool); /// Creates a `Writer` which serializes a complete V1 file index container. + /// + /// @param output_stream Destination stream for serialized index data. + /// @param pool Memory pool for writer-side allocations. + /// @return A unique pointer to a `Writer` on success. static Result> CreateWriter( const std::shared_ptr& output_stream, const std::shared_ptr& pool); @@ -126,6 +130,7 @@ class FileIndexFormat::Writer { /// Writes all column indexes. This is a terminal, one-shot operation. virtual Status WriteColumnIndexes(const FileIndexFormat::ColumnIndexes& indexes) = 0; + /// Flushes and closes the output stream supplied to `CreateWriter()`. virtual Status Close() = 0; }; diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index daa39fb53..bdb110057 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -757,6 +757,7 @@ if(PAIMON_BUILD_TESTS) core/io/vector_file_batch_reader_test.cpp core/io/data_file_meta_test.cpp core/io/data_file_index_writer_test.cpp + core/io/file_index_options_test.cpp core/io/file_index_evaluator_test.cpp core/io/single_file_writer_test.cpp core/io/rolling_blob_file_writer_test.cpp diff --git a/src/paimon/core/io/data_file_index_writer.cpp b/src/paimon/core/io/data_file_index_writer.cpp index 2d21a02b1..5c97bb3da 100644 --- a/src/paimon/core/io/data_file_index_writer.cpp +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -28,6 +28,7 @@ #include "fmt/format.h" #include "paimon/common/io/byte_array_output_stream.h" #include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/scope_guard.h" @@ -54,6 +55,10 @@ Result> DataFileIndexWriter::Create( std::vector writers; writers.reserve(options.Definitions().size()); for (const FileIndexDefinition& definition : options.Definitions()) { + if (SpecialFields::IsSystemField(definition.column_name)) { + return Status::Invalid( + fmt::format("File index column '{}' is a system field", definition.column_name)); + } int32_t field_index = logical_schema->GetFieldIndex(definition.column_name); if (field_index < 0) { return Status::Invalid( @@ -92,6 +97,9 @@ DataFileIndexWriter::DataFileIndexWriter(std::vector&& writers pool_(pool) {} Status DataFileIndexWriter::AddBatch(const std::shared_ptr& logical_batch) { + if (finished_) { + return Status::Invalid("Data file index writer has already finished"); + } for (const IndexWriterEntry& entry : writers_) { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::shared_ptr projected, @@ -108,10 +116,8 @@ Status DataFileIndexWriter::AddBatch(const std::shared_ptr& Result> DataFileIndexWriter::SerializeContainer() { FileIndexFormat::ColumnIndexes column_indexes; for (const IndexWriterEntry& entry : writers_) { - PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR serialized, + PAIMON_ASSIGN_OR_RAISE(column_indexes[entry.column_name][entry.index_type], entry.writer->SerializedBytes()); - column_indexes[entry.column_name][entry.index_type] = - std::shared_ptr(std::move(serialized)); } auto segment_output = std::make_unique( @@ -126,6 +132,10 @@ Result> DataFileIndexWriter::SerializeContainer() { } Result DataFileIndexWriter::Finish(const std::string& data_file_path) { + if (finished_) { + return Status::Invalid("Data file index writer has already finished"); + } + finished_ = true; PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, SerializeContainer()); if (static_cast(bytes->size()) <= in_manifest_threshold_) { return FileIndexWriteResult{bytes, {}}; @@ -141,7 +151,9 @@ Status DataFileIndexWriter::WriteExternal(const std::string& path, PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, file_system_->Create(path, /*overwrite=*/false)); ScopeGuard guard([this, &output]() { - [[maybe_unused]] Status _ = output->Close(); + if (output) { + [[maybe_unused]] Status _ = output->Close(); + } Abort(); }); PAIMON_ASSIGN_OR_RAISE(int64_t written, @@ -151,8 +163,9 @@ Status DataFileIndexWriter::WriteExternal(const std::string& path, path, bytes->size(), written)); } PAIMON_RETURN_NOT_OK(output->Flush()); - PAIMON_RETURN_NOT_OK(output->Close()); + Status close_status = output->Close(); output.reset(); + PAIMON_RETURN_NOT_OK(close_status); guard.Release(); return Status::OK(); } diff --git a/src/paimon/core/io/data_file_index_writer.h b/src/paimon/core/io/data_file_index_writer.h index ffb17511b..883b719fc 100644 --- a/src/paimon/core/io/data_file_index_writer.h +++ b/src/paimon/core/io/data_file_index_writer.h @@ -19,6 +19,7 @@ #pragma once +#include #include #include #include @@ -26,6 +27,7 @@ #include "paimon/core/io/file_index_options.h" #include "paimon/result.h" +#include "paimon/status.h" namespace arrow { class Field; @@ -57,6 +59,10 @@ class DataFileIndexWriter { Status AddBatch(const std::shared_ptr& logical_batch); + /// Finalizes and publishes all configured indexes. This is a terminal, one-shot operation. + /// + /// @param data_file_path Path of the data file associated with these indexes. + /// @return Embedded index bytes or the external index file name. Result Finish(const std::string& data_file_path); void Abort(); @@ -88,6 +94,7 @@ class DataFileIndexWriter { std::shared_ptr path_factory_; std::shared_ptr pool_; std::optional external_index_path_; + bool finished_ = false; }; } // namespace paimon diff --git a/src/paimon/core/io/data_file_index_writer_test.cpp b/src/paimon/core/io/data_file_index_writer_test.cpp index fc9b02e1c..e9ab7940b 100644 --- a/src/paimon/core/io/data_file_index_writer_test.cpp +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -19,6 +19,7 @@ #include "paimon/core/io/data_file_index_writer.h" +#include #include #include #include @@ -27,6 +28,7 @@ #include "arrow/ipc/json_simple.h" #include "arrow/type.h" #include "gtest/gtest.h" +#include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/core_options.h" @@ -39,9 +41,54 @@ #include "paimon/memory/bytes.h" #include "paimon/memory/memory_pool.h" #include "paimon/predicate/literal.h" +#include "paimon/testing/mock/mock_file_system.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +struct CloseFailingState { + int32_t close_count = 0; + int32_t delete_count = 0; +}; + +class CloseFailingOutputStream : public MockOutputStream { + public: + explicit CloseFailingOutputStream(const std::shared_ptr& state) + : state_(state) {} + + Result Write(const char*, int64_t size) override { + return size; + } + + Status Close() override { + ++state_->close_count; + return Status::IOError("close failed"); + } + + private: + std::shared_ptr state_; +}; + +class CloseFailingFileSystem : public MockFileSystem { + public: + explicit CloseFailingFileSystem(const std::shared_ptr& state) + : state_(state) {} + + Result> Create(const std::string&, bool) const override { + return std::unique_ptr(new CloseFailingOutputStream(state_)); + } + + Status Delete(const std::string&, bool = true) const override { + ++state_->delete_count; + return Status::OK(); + } + + private: + std::shared_ptr state_; +}; + +} // namespace class DataFileIndexWriterTest : public ::testing::Test { public: @@ -58,10 +105,8 @@ class DataFileIndexWriterTest : public ::testing::Test { Result> CreateWriter( const std::map& index_options) const { - std::map options = {{"file-system", "local"}}; - options.insert(index_options.begin(), index_options.end()); PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, - CoreOptions::FromMap(options, file_system_)); + CoreOptions::FromMap(index_options, file_system_)); PAIMON_ASSIGN_OR_RAISE(FileIndexOptions parsed, FileIndexOptions::FromCoreOptions(core_options)); return DataFileIndexWriter::Create(schema_, parsed, file_system_, path_factory_, pool_); @@ -159,4 +204,50 @@ TEST_F(DataFileIndexWriterTest, TestUnavailableWriterFailsCreation) { "do not support index writer in bloom filter"); } +TEST_F(DataFileIndexWriterTest, TestRejectSystemFieldIndex) { + std::shared_ptr key_value_schema = + SpecialFields::CompleteSequenceAndValueKindField(schema_); + for (const std::string& field_name : + {SpecialFields::SequenceNumber().Name(), SpecialFields::ValueKind().Name()}) { + ASSERT_OK_AND_ASSIGN( + CoreOptions core_options, + CoreOptions::FromMap({{"file-index.bitmap.columns", field_name}}, file_system_)); + ASSERT_OK_AND_ASSIGN(FileIndexOptions options, + FileIndexOptions::FromCoreOptions(core_options)); + ASSERT_NOK_WITH_MSG(DataFileIndexWriter::Create(key_value_schema, options, file_system_, + path_factory_, pool_), + "is a system field"); + } +} + +TEST_F(DataFileIndexWriterTest, TestFinishIsOneShot) { + ASSERT_OK_AND_ASSIGN(auto writer, + CreateWriter({{"file-index.bitmap.columns", "f0"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}})); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 10}])"))); + ASSERT_OK(writer->Finish("unused.orc")); + + ASSERT_NOK_WITH_MSG(writer->Finish("unused.orc"), "already finished"); + ASSERT_NOK_WITH_MSG(writer->AddBatch(CreateBatch(R"([{"f0": 2, "f1": 20}])")), + "already finished"); +} + +TEST_F(DataFileIndexWriterTest, TestCloseFailureClosesExternalStreamOnce) { + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, + CoreOptions::FromMap({{"file-index.bitmap.columns", "f0"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1B"}}, + file_system_)); + ASSERT_OK_AND_ASSIGN(FileIndexOptions options, FileIndexOptions::FromCoreOptions(core_options)); + auto state = std::make_shared(); + auto close_failing_file_system = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(auto writer, + DataFileIndexWriter::Create(schema_, options, close_failing_file_system, + path_factory_, pool_)); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 10}])"))); + + ASSERT_NOK_WITH_MSG(writer->Finish(path_factory_->NewPath()), "close failed"); + ASSERT_EQ(1, state->close_count); + ASSERT_EQ(1, state->delete_count); +} + } // namespace paimon::test diff --git a/src/paimon/core/io/data_file_writer.cpp b/src/paimon/core/io/data_file_writer.cpp index 5c3f4c2fe..9275fcf25 100644 --- a/src/paimon/core/io/data_file_writer.cpp +++ b/src/paimon/core/io/data_file_writer.cpp @@ -48,7 +48,7 @@ DataFileWriter::DataFileWriter( Status DataFileWriter::Write(ArrowArray* batch) { int64_t record_count = batch->length; - PAIMON_RETURN_NOT_OK(WriteRecord(batch, batch)); + PAIMON_RETURN_NOT_OK(WriteRecordWithFileIndex(batch)); seq_num_counter_->Add(record_count); return Status::OK(); } diff --git a/src/paimon/core/io/data_file_writer_base.h b/src/paimon/core/io/data_file_writer_base.h index 260402081..ccea898a7 100644 --- a/src/paimon/core/io/data_file_writer_base.h +++ b/src/paimon/core/io/data_file_writer_base.h @@ -21,6 +21,7 @@ #include #include #include +#include #include #include "arrow/c/bridge.h" @@ -32,12 +33,10 @@ #include "paimon/result.h" #include "paimon/status.h" -namespace arrow { -class Schema; -} // namespace arrow - namespace paimon { +struct KeyValueBatch; + /// Common lifecycle for data file writers which may finalize schema metadata and publish a /// file-level index. Concrete writers remain responsible for their record-specific state and /// DataFileMeta construction. @@ -84,8 +83,10 @@ class DataFileWriterBase : public SingleFileWriter converter) : Base(compression, std::move(converter)) {} - Status WriteRecord(Record record, ::ArrowArray* logical_batch) { - PAIMON_RETURN_NOT_OK(AddFileIndexBatch(logical_batch)); + /// Extracts the pre-conversion Arrow batch from record for file index construction, then + /// passes record to the underlying data file writer, which may convert it to a physical schema. + Status WriteRecordWithFileIndex(Record record) { + PAIMON_RETURN_NOT_OK(AddFileIndexBatch(GetFileIndexBatch(record))); return Base::Write(std::move(record)); } @@ -112,6 +113,16 @@ class DataFileWriterBase : public SingleFileWriter) { + return record; + } else { + static_assert(std::is_same_v, + "Unsupported data file record type"); + return record.batch.get(); + } + } + Status AddFileIndexBatch(::ArrowArray* batch) { if (!file_index_writer_) { return Status::OK(); diff --git a/src/paimon/core/io/file_index_options.cpp b/src/paimon/core/io/file_index_options.cpp index 0b170d612..a9587363a 100644 --- a/src/paimon/core/io/file_index_options.cpp +++ b/src/paimon/core/io/file_index_options.cpp @@ -19,6 +19,7 @@ #include "paimon/core/io/file_index_options.h" +#include #include #include @@ -33,6 +34,8 @@ namespace { constexpr char kFileIndexPrefix[] = "file-index."; constexpr char kColumnsSuffix[] = ".columns"; +constexpr size_t kFileIndexPrefixLength = sizeof(kFileIndexPrefix) - 1; +constexpr size_t kColumnsSuffixLength = sizeof(kColumnsSuffix) - 1; } // namespace @@ -47,20 +50,24 @@ Result FileIndexOptions::FromCoreOptions(const CoreOptions& op !StringUtils::EndsWith(key, kColumnsSuffix)) { continue; } - const size_t index_type_length = - key.size() - std::string(kFileIndexPrefix).size() - std::string(kColumnsSuffix).size(); - const std::string index_type = - key.substr(std::string(kFileIndexPrefix).size(), index_type_length); + if (key.size() < kFileIndexPrefixLength + kColumnsSuffixLength) { + return Status::Invalid(fmt::format("Invalid file index option {}", key)); + } + const size_t index_type_length = key.size() - kFileIndexPrefixLength - kColumnsSuffixLength; + const std::string index_type = key.substr(kFileIndexPrefixLength, index_type_length); if (index_type.empty()) { return Status::Invalid(fmt::format("Invalid file index option {}", key)); } + // TODO(jinli.zjw): Align malformed list option parsing (for example, "f1,f2,,") with Java. + // Update this together with ConfigParser::ParseList to keep option parsing consistent. for (std::string column_name : StringUtils::Split(value, ",", /*ignore_empty=*/false)) { StringUtils::Trim(&column_name); if (column_name.empty()) { return Status::Invalid( fmt::format("Wrong option in {}, should not have empty column", key)); } - if (column_name.find('[') != std::string::npos) { + if (column_name.find('[') != std::string::npos && + StringUtils::EndsWith(column_name, "]")) { return Status::NotImplemented( "Writing file indexes for nested map columns is not supported"); } @@ -76,8 +83,8 @@ Result FileIndexOptions::FromCoreOptions(const CoreOptions& op key == Options::FILE_INDEX_IN_MANIFEST_THRESHOLD) { continue; } - std::vector parts = StringUtils::Split( - key.substr(std::string(kFileIndexPrefix).size()), ".", /*ignore_empty=*/false); + std::vector parts = + StringUtils::Split(key.substr(kFileIndexPrefixLength), ".", /*ignore_empty=*/false); if (parts.size() != 3) { continue; } diff --git a/src/paimon/core/io/file_index_options.h b/src/paimon/core/io/file_index_options.h index ae51f05d4..7b7c019bf 100644 --- a/src/paimon/core/io/file_index_options.h +++ b/src/paimon/core/io/file_index_options.h @@ -54,6 +54,8 @@ class FileIndexOptions { } private: + FileIndexOptions() = default; + std::vector definitions_; int64_t in_manifest_threshold_ = 0; }; diff --git a/src/paimon/core/io/file_index_options_test.cpp b/src/paimon/core/io/file_index_options_test.cpp new file mode 100644 index 000000000..157203f95 --- /dev/null +++ b/src/paimon/core/io/file_index_options_test.cpp @@ -0,0 +1,58 @@ +/* + * 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/io/file_index_options.h" + +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/core/core_options.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +Result ParseOptions(const std::map& index_options) { + std::shared_ptr file_system = std::make_shared(); + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, + CoreOptions::FromMap(index_options, file_system)); + return FileIndexOptions::FromCoreOptions(core_options); +} + +} // namespace + +TEST(FileIndexOptionsTest, TestRejectOverlappingPrefixAndSuffix) { + ASSERT_NOK_WITH_MSG(ParseOptions({{"file-index.columns", "f0"}}), + "Invalid file index option file-index.columns"); +} + +TEST(FileIndexOptionsTest, TestNestedMapColumnSyntax) { + ASSERT_OK_AND_ASSIGN(FileIndexOptions options, + ParseOptions({{"file-index.bitmap.columns", "col[key"}})); + ASSERT_EQ(1, options.Definitions().size()); + ASSERT_EQ("col[key", options.Definitions()[0].column_name); + + ASSERT_NOK_WITH_MSG(ParseOptions({{"file-index.bitmap.columns", "col[key]"}}), + "nested map columns is not supported"); +} + +} // namespace paimon::test diff --git a/src/paimon/core/io/key_value_data_file_writer.cpp b/src/paimon/core/io/key_value_data_file_writer.cpp index 56a8bcd4b..9c32e0674 100644 --- a/src/paimon/core/io/key_value_data_file_writer.cpp +++ b/src/paimon/core/io/key_value_data_file_writer.cpp @@ -75,8 +75,7 @@ Status KeyValueDataFileWriter::Write(KeyValueBatch batch) { // update delete row count delete_row_count_ += batch.delete_row_count; - ::ArrowArray* logical_batch = batch.batch.get(); - return WriteRecord(std::move(batch), logical_batch); + return WriteRecordWithFileIndex(std::move(batch)); } Result> KeyValueDataFileWriter::GetResult() {