feat(mergetree): add compact manager, task abstractions and file rewrite task (#105)

Author: lucasfang

src/paimon/core/mergetree/compact/file_rewrite_compact_task.h

@@ -0,0 +1,67 @@
+/*
+ * 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 <cstdint>
+#include <memory>
+#include <vector>
+
+#include "paimon/core/compact/compact_task.h"
+#include "paimon/core/compact/compact_unit.h"
+#include "paimon/core/mergetree/compact/compact_rewriter.h"
+#include "paimon/core/mergetree/sorted_run.h"
+
+namespace paimon {
+
+/// Compact task for file rewrite compaction.
+class FileRewriteCompactTask : public CompactTask {
+ public:
+    FileRewriteCompactTask(const std::shared_ptr<CompactRewriter>& rewriter,
+                           const CompactUnit& unit, bool drop_delete,
+                           const std::shared_ptr<CompactionMetrics::Reporter>& metrics_reporter)
+        : CompactTask(metrics_reporter),
+          rewriter_(rewriter),
+          output_level_(unit.output_level),
+          files_(unit.files),
+          drop_delete_(drop_delete) {}
+
+ protected:
+    Result<std::shared_ptr<CompactResult>> DoCompact() override {
+        auto result = std::make_shared<CompactResult>();
+        for (const auto& file : files_) {
+            PAIMON_RETURN_NOT_OK(RewriteFile(file, result.get()));
+        }
+        return result;
+    }
+
+ private:
+    Status RewriteFile(const std::shared_ptr<DataFileMeta>& file, CompactResult* to_update) {
+        std::vector<std::vector<SortedRun>> candidate = {{SortedRun::FromSingle(file)}};
+        PAIMON_ASSIGN_OR_RAISE(CompactResult rewritten,
+                               rewriter_->Rewrite(output_level_, drop_delete_, candidate));
+        return to_update->Merge(rewritten);
+    }
+
+    std::shared_ptr<CompactRewriter> rewriter_;
+    int32_t output_level_;
+    std::vector<std::shared_ptr<DataFileMeta>> files_;
+    bool drop_delete_;
+};
+
+}  // namespace paimon

src/paimon/core/mergetree/compact/file_rewrite_compact_task_test.cpp

@@ -0,0 +1,167 @@
+/*
+ * 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/mergetree/compact/file_rewrite_compact_task.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paimon/core/io/data_file_meta.h"
+#include "paimon/core/manifest/file_source.h"
+#include "paimon/core/stats/simple_stats.h"
+#include "paimon/testing/utils/testharness.h"
+
+namespace paimon::test {
+namespace {
+
+class SpyCompactRewriter final : public CompactRewriter {
+ public:
+    Result<CompactResult> Rewrite(int32_t output_level, bool drop_delete,
+                                  const std::vector<std::vector<SortedRun>>& sections) override {
+        output_levels.push_back(output_level);
+        drop_deletes.push_back(drop_delete);
+
+        EXPECT_EQ(sections.size(), 1);
+        EXPECT_EQ(sections[0].size(), 1);
+        EXPECT_EQ(sections[0][0].Files().size(), 1);
+
+        auto input = sections[0][0].Files()[0];
+        rewritten_files.push_back(input);
+
+        PAIMON_ASSIGN_OR_RAISE(auto upgraded, input->Upgrade(output_level));
+        return CompactResult({input}, {upgraded});
+    }
+
+    Result<CompactResult> Upgrade(int32_t, const std::shared_ptr<DataFileMeta>&) override {
+        return Status::Invalid("Upgrade should not be called by FileRewriteCompactTask");
+    }
+
+    Status Close() override {
+        return Status::OK();
+    }
+
+    std::vector<int32_t> output_levels;
+    std::vector<bool> drop_deletes;
+    std::vector<std::shared_ptr<DataFileMeta>> rewritten_files;
+};
+
+class FailingCompactRewriter final : public CompactRewriter {
+ public:
+    explicit FailingCompactRewriter(size_t fail_on_call) : fail_on_call_(fail_on_call) {}
+
+    Result<CompactResult> Rewrite(int32_t, bool,
+                                  const std::vector<std::vector<SortedRun>>& sections) override {
+        EXPECT_EQ(sections.size(), 1);
+        EXPECT_EQ(sections[0].size(), 1);
+        EXPECT_EQ(sections[0][0].Files().size(), 1);
+
+        rewritten_files.push_back(sections[0][0].Files()[0]);
+        ++call_count_;
+        if (call_count_ == fail_on_call_) {
+            return Status::Invalid("injected rewrite failure");
+        }
+
+        auto input = sections[0][0].Files()[0];
+        PAIMON_ASSIGN_OR_RAISE(auto upgraded, input->Upgrade(/*new_level=*/1));
+        return CompactResult({input}, {upgraded});
+    }
+
+    Result<CompactResult> Upgrade(int32_t, const std::shared_ptr<DataFileMeta>&) override {
+        return Status::Invalid("Upgrade should not be called by FileRewriteCompactTask");
+    }
+
+    Status Close() override {
+        return Status::OK();
+    }
+
+    size_t call_count() const {
+        return call_count_;
+    }
+
+    std::vector<std::shared_ptr<DataFileMeta>> rewritten_files;
+
+ private:
+    size_t fail_on_call_;
+    size_t call_count_ = 0;
+};
+
+std::shared_ptr<DataFileMeta> NewAppendFile(const std::string& file_name, int64_t row_count,
+                                            int64_t min_sequence_number,
+                                            int64_t max_sequence_number) {
+    return DataFileMeta::ForAppend(file_name, /*file_size=*/row_count, row_count,
+                                   SimpleStats::EmptyStats(), min_sequence_number,
+                                   max_sequence_number, /*schema_id=*/0, FileSource::Append(),
+                                   std::nullopt, std::nullopt, std::nullopt, std::nullopt)
+        .value();
+}
+
+}  // namespace
+
+TEST(FileRewriteCompactTaskTest, TestRewriteEachFileAndMergeResults) {
+    auto file1 = NewAppendFile("file-1", /*row_count=*/10, /*min_sequence_number=*/0,
+                               /*max_sequence_number=*/9);
+    auto file2 = NewAppendFile("file-2", /*row_count=*/20, /*min_sequence_number=*/10,
+                               /*max_sequence_number=*/29);
+    CompactUnit unit(/*output_level=*/2, {file1, file2}, /*file_rewrite=*/true);
+
+    auto rewriter = std::make_shared<SpyCompactRewriter>();
+    FileRewriteCompactTask task(rewriter, unit, /*drop_delete=*/true, /*metrics_reporter=*/nullptr);
+
+    ASSERT_OK_AND_ASSIGN(std::shared_ptr<CompactResult> result, task.Execute());
+
+    ASSERT_EQ(rewriter->output_levels, std::vector<int32_t>({2, 2}));
+    ASSERT_EQ(rewriter->drop_deletes, std::vector<bool>({true, true}));
+    ASSERT_EQ(rewriter->rewritten_files.size(), 2);
+    EXPECT_EQ(rewriter->rewritten_files[0], file1);
+    EXPECT_EQ(rewriter->rewritten_files[1], file2);
+
+    ASSERT_EQ(result->Before().size(), 2);
+    EXPECT_EQ(result->Before()[0], file1);
+    EXPECT_EQ(result->Before()[1], file2);
+
+    ASSERT_EQ(result->After().size(), 2);
+    EXPECT_EQ(result->After()[0]->file_name, file1->file_name);
+    EXPECT_EQ(result->After()[1]->file_name, file2->file_name);
+    EXPECT_EQ(result->After()[0]->level, 2);
+    EXPECT_EQ(result->After()[1]->level, 2);
+}
+
+TEST(FileRewriteCompactTaskTest, TestStopOnRewriteFailure) {
+    auto file1 = NewAppendFile("file-1", /*row_count=*/10, /*min_sequence_number=*/0,
+                               /*max_sequence_number=*/9);
+    auto file2 = NewAppendFile("file-2", /*row_count=*/20, /*min_sequence_number=*/10,
+                               /*max_sequence_number=*/29);
+    auto file3 = NewAppendFile("file-3", /*row_count=*/30, /*min_sequence_number=*/30,
+                               /*max_sequence_number=*/59);
+    CompactUnit unit(/*output_level=*/2, {file1, file2, file3}, /*file_rewrite=*/true);
+
+    auto rewriter = std::make_shared<FailingCompactRewriter>(/*fail_on_call=*/2);
+    FileRewriteCompactTask task(rewriter, unit, /*drop_delete=*/false,
+                                /*metrics_reporter=*/nullptr);
+
+    ASSERT_NOK_WITH_MSG(task.Execute(), "injected rewrite failure");
+    ASSERT_EQ(rewriter->call_count(), 2);
+    ASSERT_EQ(rewriter->rewritten_files.size(), 2);
+    EXPECT_EQ(rewriter->rewritten_files[0], file1);
+    EXPECT_EQ(rewriter->rewritten_files[1], file2);
+}
+
+}  // namespace paimon::test