feat(core): complete file index write lifecycle (#279)

Author: zjw1111

src/paimon/CMakeLists.txt

@@ -779,6 +779,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_writer_base_test.cpp
                     core/io/data_file_index_writer_test.cpp
                     core/io/file_index_options_test.cpp
                     core/io/file_index_evaluator_test.cpp

src/paimon/core/append/append_only_writer.cpp

@@ -284,7 +284,9 @@ Status AppendOnlyWriter::Close() {
     for (const auto& file : compact_after_) {
         // AppendOnlyCompactManager will rewrite the file and no file upgrade will occur, so we
         // can directly delete the file in compact_after_.
-        [[maybe_unused]] auto s = fs->Delete(path_factory_->ToPath(file));
+        for (const std::string& path : path_factory_->CollectFiles(file)) {
+            [[maybe_unused]] Status s = fs->Delete(path);
+        }
     }
 
     if (writer_) {

src/paimon/core/append/append_only_writer_test.cpp

@@ -443,6 +443,9 @@ TEST_F(AppendOnlyWriterTest, TestWriteAndClose) {
     std::map<std::string, std::string> raw_options;
     raw_options[Options::FILE_FORMAT] = "orc";
     raw_options[Options::FILE_SYSTEM] = "local";
+    raw_options[Options::TARGET_FILE_ROW_NUM] = "1";
+    raw_options["file-index.bitmap.columns"] = "f0";
+    raw_options[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1B";
     ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(raw_options));
 
     arrow::FieldVector fields = {arrow::field("f0", arrow::utf8())};
@@ -475,11 +478,15 @@ TEST_F(AppendOnlyWriterTest, TestWriteAndClose) {
     ASSERT_OK_AND_ASSIGN(auto record_batch, batch_builder.Finish());
     ASSERT_OK(writer->Write(std::move(record_batch)));
     ASSERT_TRUE(ArrowArrayIsReleased(&arrow_array));
-    ASSERT_OK(writer->Close());
 
     auto file_system = std::make_shared<LocalFileSystem>();
     std::vector<BasicFileStatus> file_status_list;
     ASSERT_OK(file_system->ListDir(dir->Str(), &file_status_list));
+    ASSERT_EQ(2, file_status_list.size());
+
+    ASSERT_OK(writer->Close());
+    file_status_list.clear();
+    ASSERT_OK(file_system->ListDir(dir->Str(), &file_status_list));
     ASSERT_TRUE(file_status_list.empty());
 }
 
@@ -634,9 +641,13 @@ TEST_F(AppendOnlyWriterTest, TestCloseDeletesCompactAfterFiles) {
     auto compact_manager = std::make_shared<FakeCompactManager>();
 
     auto compact_after = NewAppendFile("compact-after.orc", 1, 0, 0);
+    compact_after->extra_files = {"compact-after.orc.index"};
     auto compact_after_path = path_factory->ToPath(compact_after->file_name);
+    auto compact_after_index_path = path_factory->ToPath(compact_after->extra_files[0].value());
     ASSERT_OK_AND_ASSIGN(auto output, options.GetFileSystem()->Create(compact_after_path, true));
     ASSERT_OK(output->Close());
+    ASSERT_OK_AND_ASSIGN(output, options.GetFileSystem()->Create(compact_after_index_path, true));
+    ASSERT_OK(output->Close());
 
     auto result =
         std::make_shared<CompactResult>(std::vector<std::shared_ptr<DataFileMeta>>{},
@@ -653,8 +664,10 @@ TEST_F(AppendOnlyWriterTest, TestCloseDeletesCompactAfterFiles) {
 
     ASSERT_OK(writer->Sync());
     ASSERT_TRUE(options.GetFileSystem()->Exists(compact_after_path).value());
+    ASSERT_TRUE(options.GetFileSystem()->Exists(compact_after_index_path).value());
     ASSERT_OK(writer->Close());
     ASSERT_FALSE(options.GetFileSystem()->Exists(compact_after_path).value());
+    ASSERT_FALSE(options.GetFileSystem()->Exists(compact_after_index_path).value());
     ASSERT_TRUE(compact_manager->request_cancel_called);
     ASSERT_TRUE(compact_manager->wait_called);
     ASSERT_TRUE(compact_manager->close_called);

src/paimon/core/io/data_file_index_writer_test.cpp

@@ -143,9 +143,10 @@ class DataFileIndexWriterTest : public ::testing::Test {
     std::shared_ptr<arrow::Schema> schema_;
 };
 
-TEST_F(DataFileIndexWriterTest, TestBitmapAndRangeBitmapEmbeddedRoundTrip) {
+TEST_F(DataFileIndexWriterTest, TestMultipleIndexesOnSameColumnEmbeddedRoundTrip) {
     ASSERT_OK_AND_ASSIGN(auto writer,
                          CreateWriter({{"file-index.bitmap.columns", "f0"},
+                                       {"file-index.bsi.columns", "f0"},
                                        {"file-index.range-bitmap.columns", "f1"},
                                        {"file-index.range-bitmap.f1.chunk-size", "1KB"},
                                        {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}}));
@@ -159,12 +160,14 @@ TEST_F(DataFileIndexWriterTest, TestBitmapAndRangeBitmapEmbeddedRoundTrip) {
     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 f0_readers, ReadColumn(reader.get(), "f0"));
+    ASSERT_EQ(2, f0_readers.size());
+    for (const std::shared_ptr<FileIndexReader>& f0_reader : f0_readers) {
+        ASSERT_OK_AND_ASSIGN(auto equal_result, f0_reader->VisitEqual(Literal(1)));
+        ASSERT_EQ("{0,2}", equal_result->ToString());
+        ASSERT_OK_AND_ASSIGN(auto null_result, f0_reader->VisitIsNull());
+        ASSERT_EQ("{3}", null_result->ToString());
+    }
 
     ASSERT_OK_AND_ASSIGN(auto range_readers, ReadColumn(reader.get(), "f1"));
     ASSERT_EQ(1, range_readers.size());

src/paimon/core/io/data_file_writer_base.h

@@ -88,6 +88,7 @@ class DataFileWriterBase : public SingleFileWriter<Record, std::shared_ptr<DataF
     /// 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(this->CheckNotClosed());
         PAIMON_RETURN_NOT_OK(AddFileIndexBatch(GetFileIndexBatch(record)));
         return Base::Write(std::move(record));
     }

src/paimon/core/io/data_file_writer_base_test.cpp

@@ -0,0 +1,97 @@
+/*
+ * 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_writer_base.h"
+
+#include <memory>
+#include <utility>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "arrow/ipc/json_simple.h"
+#include "gtest/gtest.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/defs.h"
+#include "paimon/format/file_format.h"
+#include "paimon/memory/memory_pool.h"
+#include "paimon/testing/utils/testharness.h"
+
+namespace paimon::test {
+
+class IndexedDataFileWriter : public DataFileWriterBase<::ArrowArray*> {
+ public:
+    IndexedDataFileWriter() : DataFileWriterBase(/*compression=*/"zstd", /*converter=*/nullptr) {}
+
+    Status Write(::ArrowArray* batch) override {
+        return WriteRecordWithFileIndex(batch);
+    }
+
+    Result<std::shared_ptr<DataFileMeta>> GetResult() override {
+        return std::shared_ptr<DataFileMeta>();
+    }
+};
+
+TEST(DataFileWriterBaseTest, WriteAfterCloseDoesNotConsumeIndexBatch) {
+    auto dir = UniqueTestDirectory::Create();
+    ASSERT_TRUE(dir);
+    std::shared_ptr<MemoryPool> pool = GetDefaultPool();
+    std::shared_ptr<arrow::Schema> schema = arrow::schema({arrow::field("col", arrow::int32())});
+    std::shared_ptr<arrow::DataType> data_type = arrow::struct_(schema->fields());
+    ASSERT_OK_AND_ASSIGN(
+        CoreOptions options,
+        CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"},
+                              {"file-index.bitmap.columns", "col"},
+                              {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}}));
+    std::shared_ptr<FileSystem> file_system = options.GetFileSystem();
+    auto path_factory = std::make_shared<DataFilePathFactory>();
+    ASSERT_OK(path_factory->Init(dir->Str(), "orc", "data-", nullptr));
+    ASSERT_OK_AND_ASSIGN(FileIndexOptions file_index_options,
+                         FileIndexOptions::FromCoreOptions(options));
+    ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<DataFileIndexWriter> file_index_writer,
+        DataFileIndexWriter::Create(schema, file_index_options, file_system, path_factory, pool));
+
+    ArrowSchema c_schema;
+    ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok());
+    ASSERT_OK_AND_ASSIGN(std::shared_ptr<WriterBuilder> writer_builder,
+                         options.GetFileFormat()->CreateWriterBuilder(&c_schema,
+                                                                      /*batch_size=*/100));
+    IndexedDataFileWriter writer;
+    writer.SetFileIndexWriter(std::move(file_index_writer), schema);
+    ASSERT_OK(writer.Init(file_system, path_factory->NewPath(), writer_builder));
+
+    std::shared_ptr<arrow::Array> first_array =
+        arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([[1]])").ValueOrDie();
+    ::ArrowArray first_batch;
+    ASSERT_TRUE(arrow::ExportArray(*first_array, &first_batch).ok());
+    ASSERT_OK(writer.Write(&first_batch));
+    ASSERT_OK(writer.Close());
+
+    std::shared_ptr<arrow::Array> second_array =
+        arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([[2]])").ValueOrDie();
+    ::ArrowArray second_batch;
+    ASSERT_TRUE(arrow::ExportArray(*second_array, &second_batch).ok());
+    ASSERT_NOK_WITH_MSG(writer.Write(&second_batch), "Writer has already closed");
+    ASSERT_NE(nullptr, second_batch.release);
+    ArrowArrayRelease(&second_batch);
+}
+
+}  // namespace paimon::test

src/paimon/core/io/key_value_data_file_writer_factory.cpp

@@ -62,6 +62,7 @@ KeyValueDataFileWriterFactory::CreateWriter() const {
     auto writer = std::make_unique<KeyValueDataFileWriter>(
         GetFileCompression(), std::move(converter), schema_id_, level_, file_source_, primary_keys_,
         resources.stats_extractor, write_schema_, path_factory_->IsExternalPath(), pool_);
+    // Changelog files are consumed sequentially and intentionally do not produce file indexes.
     if (!is_changelog_) {
         PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<DataFileIndexWriter> file_index_writer,
                                CreateFileIndexWriter(write_schema_, path_factory_));

src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp

@@ -90,6 +90,7 @@ ShreddingKeyValueDataFileWriterFactory::CreateShreddedWriter(
     auto writer = std::make_unique<KeyValueDataFileWriter>(
         compression, std::move(batch_converter), schema_id_, level_, file_source_, primary_keys_,
         resources.stats_extractor, file_schema, path_factory_->IsExternalPath(), pool_);
+    // Changelog files are consumed sequentially and intentionally do not produce file indexes.
     if (!is_changelog_) {
         PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<DataFileIndexWriter> file_index_writer,
                                CreateFileIndexWriter(write_schema_, path_factory_));

src/paimon/core/io/single_file_writer.h

@@ -133,6 +133,14 @@ class SingleFileWriter : public FileWriter<T, R> {
     }
 
  protected:
+    /// Returns an error if this writer has already been closed.
+    Status CheckNotClosed() const {
+        if (PAIMON_UNLIKELY(closed_)) {
+            return Status::Invalid("Writer has already closed!");
+        }
+        return Status::OK();
+    }
+
     /// Hook called after Flush() and before Finish() during Close().
     /// Subclasses can override to update per-field metadata before the file is finalized.
     virtual Status BeforeFinish() {
@@ -196,9 +204,7 @@ Status SingleFileWriter<T, R>::Init(const std::shared_ptr<FileSystem>& fs, const
 
 template <typename T, typename R>
 Status SingleFileWriter<T, R>::Write(T record) {
-    if (PAIMON_UNLIKELY(closed_)) {
-        return Status::Invalid("Writer has already closed!");
-    }
+    PAIMON_RETURN_NOT_OK(CheckNotClosed());
     ScopeGuard guard([this]() -> void { this->Abort(); });
     int64_t record_count = 0;
     if (!converter_) {

src/paimon/core/mergetree/lookup_levels.h

@@ -45,6 +45,8 @@ struct RemoteSstFile {
 template <typename T>
 class LookupLevels : public Levels::DropFileCallback {
  public:
+    static constexpr const char* REMOTE_LOOKUP_FILE_SUFFIX = ".lookup";
+
     static Result<std::unique_ptr<LookupLevels<T>>> Create(
         const std::shared_ptr<FileSystem>& fs, const BinaryRow& partition, int32_t bucket,
         const CoreOptions& options, const std::shared_ptr<SchemaManager>& schema_manager,
@@ -119,8 +121,6 @@ class LookupLevels : public Levels::DropFileCallback {
     Result<std::shared_ptr<PersistProcessor<T>>> GetOrCreateProcessor(
         int64_t schema_id, const std::string& ser_version);
 
-    static constexpr const char* REMOTE_LOOKUP_FILE_SUFFIX = ".lookup";
-
  private:
     std::shared_ptr<MemoryPool> pool_;
     std::shared_ptr<FileSystem> fs_;