diff --git a/include/paimon/defs.h b/include/paimon/defs.h index d1ebf507..e944587f 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 b46dee8c..3993b624 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,9 +91,15 @@ class MemoryPool; 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 @@ -98,18 +107,38 @@ 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. + /// + /// @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); + 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; + + /// Flushes and closes the output stream supplied to `CreateWriter()`. + virtual Status Close() = 0; +}; + /// Reader for file index file. 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/CMakeLists.txt b/src/paimon/CMakeLists.txt index 9b0807b6..bdb11005 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,8 @@ 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_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/common/defs.cpp b/src/paimon/common/defs.cpp index bac4f16f..ef35940e 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 85500845..fab5c7a7 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,128 @@ 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"); + } + + 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")); + 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")); + // 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(), "index body", &body_length)); + } + } + } + + head_length += 4; // The trailing redundant-length field(4). + PAIMON_RETURN_NOT_OK( + 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)); + // 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 = head_length; + for (const auto& [column_name, column_indexes] : indexes) { + // Write column name. + PAIMON_RETURN_NOT_OK(data_output.WriteString(column_name)); + // 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)); + // 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)); + } + } + } + // Write redundant length for future format extensions. + return data_output.WriteValue(kRedundantLength); + } + + template + 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); + } + + std::shared_ptr output_stream_; + bool written_ = false; + bool closed_ = false; +}; + class FileIndexFormatReaderImpl : public FileIndexFormat::Reader { public: using HeaderType = @@ -153,4 +277,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 7851d57e..40989f7e 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,14 +58,25 @@ class FileIndexFormatTest : public ::testing::Test { std::shared_ptr pool_; }; -TEST_F(FileIndexFormatTest, TestCreateEmptyFileIndexReader) { +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 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(pool_.get())); + + ASSERT_EQ(expected, std::vector(actual->data(), actual->data() + actual->size())); 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 new file mode 100644 index 00000000..bc6d1f1a --- /dev/null +++ b/src/paimon/common/io/byte_array_output_stream.cpp @@ -0,0 +1,80 @@ +/* + * 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 +#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(std::unique_ptr&& output) + : output_(std::move(output)) { + assert(output_); +} + +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; + } + return size; +} + +Status ByteArrayOutputStream::Close() { + closed_ = true; + return Status::OK(); +} + +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. + 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(size)); + 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 00000000..9b87ca42 --- /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: + /// Takes ownership of an initialized segmented output stream. + explicit ByteArrayOutputStream(std::unique_ptr&& output); + + ~ByteArrayOutputStream() override = default; + + Result Write(const char* buffer, int64_t size) override; + + Status Flush() override { + return Status::OK(); + } + + Result GetPos() const override { + return output_->CurrentSize(); + } + + Result GetUri() const override { + return std::string(); + } + + Status Close() override; + + /// Closes the stream and returns its contents as an exactly-sized contiguous byte array. + /// @note The caller must keep `pool` alive until the returned bytes are destroyed. + Result> Finish(MemoryPool* pool); + + private: + std::unique_ptr output_; + std::shared_ptr result_; + 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 00000000..bd185095 --- /dev/null +++ b/src/paimon/common/io/byte_array_output_stream_test.cpp @@ -0,0 +1,86 @@ +/* + * 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 "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(); + auto output = std::make_unique(/*segment_size=*/2, pool); + std::shared_ptr stream = + 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); + 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(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(pool.get())); + ASSERT_EQ(result, repeated); + stream.reset(); + ASSERT_EQ(6, pool->CurrentUsage()); +} + +TEST(ByteArrayOutputStreamTest, TestWriteValidation) { + 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(pool.get())); + ASSERT_EQ(0, result->size()); +} + +TEST(ByteArrayOutputStreamTest, TestCallerKeepsMemoryPoolAlive) { + std::shared_ptr pool = GetMemoryPool(); + auto output = std::make_unique(/*segment_size=*/8, pool); + std::shared_ptr stream = + 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(pool.get())); + + stream.reset(); + ASSERT_GT(pool->CurrentUsage(), 0); + ASSERT_EQ("data", std::string(result->data(), result->size())); + + result.reset(); + 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 4e606370..0a5dd574 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 5355f72b..2d0d274a 100644 --- a/src/paimon/common/io/memory_segment_output_stream.cpp +++ b/src/paimon/common/io/memory_segment_output_stream.cpp @@ -54,11 +54,7 @@ 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); - } - auto segment = MemorySegment::Wrap(bytes); + MemorySegment segment = MemorySegment::WrapView(data, size); Write(segment, 0, segment.Size()); } 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 61fbfe30..69c207ce 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 8ef6e135..11401609 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 578950fb..1c2e164b 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 9eb28988..53ef4ad0 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 ebc127ed..0054a5b5 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 e12374dd..d0b677a7 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 00000000..5c97bb3d --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -0,0 +1,179 @@ +/* + * 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/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" +#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()) { + 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( + 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) { + 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, + 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 (const IndexWriterEntry& entry : writers_) { + PAIMON_ASSIGN_OR_RAISE(column_indexes[entry.column_name][entry.index_type], + entry.writer->SerializedBytes()); + } + + 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(pool_.get()); +} + +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, {}}; + } + + 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]() { + if (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()); + Status close_status = output->Close(); + output.reset(); + PAIMON_RETURN_NOT_OK(close_status); + 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 00000000..883b719f --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer.h @@ -0,0 +1,100 @@ +/* + * 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 + +#include "paimon/core/io/file_index_options.h" +#include "paimon/result.h" +#include "paimon/status.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); + + /// 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(); + + 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_; + 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 new file mode 100644 index 00000000..e9ab7940 --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -0,0 +1,253 @@ +/* + * 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 + +#include "arrow/c/bridge.h" +#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" +#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/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: + 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 { + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, + 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_); + } + + 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 checked_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"); +} + +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 4ed3e040..9275fcf2 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(WriteRecordWithFileIndex(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 60cc808a..f56f3495 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 00000000..ccea898a --- /dev/null +++ b/src/paimon/core/io/data_file_writer_base.h @@ -0,0 +1,145 @@ +/* + * 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 + +#include "arrow/c/bridge.h" +#include "arrow/type.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" +#include "paimon/result.h" +#include "paimon/status.h" + +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. +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)) {} + + /// 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)); + } + + 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: + static ::ArrowArray* GetFileIndexBatch(Record& record) { + if constexpr (std::is_same_v) { + 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(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr logical_array, + arrow::ImportArray(batch, logical_type_)); + std::shared_ptr logical_batch = + 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 b929dde8..07195ab7 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 c727b47d..cab942f0 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 00000000..a9587363 --- /dev/null +++ b/src/paimon/core/io/file_index_options.cpp @@ -0,0 +1,109 @@ +/* + * 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 "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"; +constexpr size_t kFileIndexPrefixLength = sizeof(kFileIndexPrefix) - 1; +constexpr size_t kColumnsSuffixLength = sizeof(kColumnsSuffix) - 1; + +} // 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; + } + 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 && + StringUtils::EndsWith(column_name, "]")) { + 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(kFileIndexPrefixLength), ".", /*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 00000000..7b7c019b --- /dev/null +++ b/src/paimon/core/io/file_index_options.h @@ -0,0 +1,63 @@ +/* + * 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: + FileIndexOptions() = default; + + std::vector definitions_; + int64_t in_manifest_threshold_ = 0; +}; + +} // namespace paimon 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 00000000..157203f9 --- /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 9393c7c3..9c32e067 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,7 @@ 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(); + return WriteRecordWithFileIndex(std::move(batch)); } Result> KeyValueDataFileWriter::GetResult() { @@ -120,14 +103,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 e1e3fd92..eb7a2efc 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 07d50b98..8f388559 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 6e4843bb..0e4e8219 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 8ac583ee..30d4c9fc 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 99507b57..6db3a699 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_()); } diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp index 82eb5257..5916fc97 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()),