| /* |
| * 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 <map> |
| #include <memory> |
| #include <string> |
| #include <vector> |
| |
| #include "arrow/c/bridge.h" |
| #include "gtest/gtest.h" |
| #include "paimon/commit_context.h" |
| #include "paimon/common/data/binary_row.h" |
| #include "paimon/common/factories/io_hook.h" |
| #include "paimon/common/utils/scope_guard.h" |
| #include "paimon/core/append/bucketed_append_compact_manager.h" |
| #include "paimon/core/io/data_file_meta.h" |
| #include "paimon/core/operation/append_only_file_store_write.h" |
| #include "paimon/core/operation/restore_files.h" |
| #include "paimon/core/table/sink/commit_message_impl.h" |
| #include "paimon/core/table/source/data_split_impl.h" |
| #include "paimon/executor.h" |
| #include "paimon/file_store_commit.h" |
| #include "paimon/file_store_write.h" |
| #include "paimon/result.h" |
| #include "paimon/testing/utils/binary_row_generator.h" |
| #include "paimon/testing/utils/data_generator.h" |
| #include "paimon/testing/utils/io_exception_helper.h" |
| #include "paimon/testing/utils/test_helper.h" |
| #include "paimon/testing/utils/testharness.h" |
| #include "paimon/write_context.h" |
| |
| namespace paimon::test { |
| |
| class AppendCompactionInteTest : public testing::Test, |
| public ::testing::WithParamInterface<std::string> { |
| public: |
| void SetUp() override { |
| pool_ = GetDefaultPool(); |
| } |
| |
| void PrepareSimpleAppendData(const std::shared_ptr<DataGenerator>& gen, bool with_dv, |
| TestHelper* helper, int64_t* identifier) { |
| auto& commit_identifier = *identifier; |
| std::vector<BinaryRow> datas_1; |
| datas_1.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Alice"), 10, 1, 11.1}, pool_.get())); |
| datas_1.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Bob"), 10, 0, 12.1}, pool_.get())); |
| datas_1.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Emily"), 10, 0, 13.1}, pool_.get())); |
| datas_1.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Tony"), 10, 0, 14.1}, pool_.get())); |
| datas_1.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Lucy"), 20, 1, 14.1}, pool_.get())); |
| ASSERT_OK_AND_ASSIGN(auto batches_1, gen->SplitArrayByPartitionAndBucket(datas_1)); |
| ASSERT_EQ(3, batches_1.size()); |
| ASSERT_OK_AND_ASSIGN( |
| auto commit_msgs, |
| helper->WriteAndCommit(std::move(batches_1), commit_identifier++, std::nullopt)); |
| ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot1, helper->LatestSnapshot()); |
| ASSERT_TRUE(snapshot1); |
| ASSERT_EQ(1, snapshot1.value().Id()); |
| ASSERT_EQ(5, snapshot1.value().TotalRecordCount().value()); |
| ASSERT_EQ(5, snapshot1.value().DeltaRecordCount().value()); |
| |
| std::vector<BinaryRow> datas_2; |
| datas_2.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Emily"), 10, 0, 15.1}, pool_.get())); |
| datas_2.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Bob"), 10, 0, 12.1}, pool_.get())); |
| datas_2.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Alex"), 10, 0, 16.1}, pool_.get())); |
| datas_2.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Paul"), 20, 1, NullType()}, pool_.get())); |
| ASSERT_OK_AND_ASSIGN(auto batches_2, gen->SplitArrayByPartitionAndBucket(datas_2)); |
| ASSERT_EQ(2, batches_2.size()); |
| ASSERT_OK_AND_ASSIGN( |
| auto commit_msgs_2, |
| helper->WriteAndCommit(std::move(batches_2), commit_identifier++, std::nullopt)); |
| ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot2, helper->LatestSnapshot()); |
| ASSERT_TRUE(snapshot2); |
| ASSERT_EQ(2, snapshot2.value().Id()); |
| ASSERT_EQ(9, snapshot2.value().TotalRecordCount().value()); |
| ASSERT_EQ(4, snapshot2.value().DeltaRecordCount().value()); |
| |
| std::vector<BinaryRow> datas_3; |
| datas_3.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("David"), 10, 0, 17.1}, pool_.get())); |
| ASSERT_OK_AND_ASSIGN(auto batches_3, gen->SplitArrayByPartitionAndBucket(datas_3)); |
| ASSERT_EQ(1, batches_3.size()); |
| ASSERT_OK_AND_ASSIGN( |
| auto commit_msgs_3, |
| helper->WriteAndCommit(std::move(batches_3), commit_identifier++, std::nullopt)); |
| ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot3, helper->LatestSnapshot()); |
| ASSERT_TRUE(snapshot3); |
| ASSERT_EQ(3, snapshot3.value().Id()); |
| ASSERT_EQ(10, snapshot3.value().TotalRecordCount().value()); |
| ASSERT_EQ(1, snapshot3.value().DeltaRecordCount().value()); |
| |
| // @note: for append-only tables in Spark, native row-level deletes aren't supported during |
| // writing. Instead, deletions are expressed by committing a Deletion Vector (DV) file |
| // externally. |
| if (with_dv) { |
| auto partition = BinaryRowGenerator::GenerateRow({10}, pool_.get()); |
| int32_t bucket = 1; |
| auto abstract_write = dynamic_cast<AbstractFileStoreWrite*>(helper->write_.get()); |
| ASSERT_OK_AND_ASSIGN(auto restore_files, |
| abstract_write->ScanExistingFileMetas(partition, bucket)); |
| ASSERT_OK_AND_ASSIGN( |
| auto dv_maintainer, |
| abstract_write->dv_maintainer_factory_->Create( |
| partition, bucket, std::vector<std::shared_ptr<IndexFileMeta>>{})); |
| for (const auto& data_file : restore_files->DataFiles()) { |
| ASSERT_OK(dv_maintainer->NotifyNewDeletion(data_file->file_name, 0)); |
| } |
| ASSERT_OK_AND_ASSIGN(auto index_file_meta, dv_maintainer->WriteDeletionVectorsIndex()); |
| |
| auto commit_message = std::make_shared<CommitMessageImpl>( |
| partition, bucket, 2, DataIncrement({}, {}, {}, {index_file_meta.value()}, {}), |
| CompactIncrement({}, {}, {})); |
| std::vector<std::shared_ptr<CommitMessage>> commit_messages; |
| commit_messages.push_back(commit_message); |
| ASSERT_OK(helper->commit_->Commit(commit_messages, commit_identifier++)); |
| } |
| } |
| |
| private: |
| std::shared_ptr<MemoryPool> pool_; |
| }; |
| |
| std::vector<std::string> GetTestValuesForAppendCompactionInteTest() { |
| std::vector<std::string> values; |
| values.emplace_back("parquet"); |
| #ifdef PAIMON_ENABLE_ORC |
| values.emplace_back("orc"); |
| #endif |
| #ifdef PAIMON_ENABLE_AVRO |
| values.emplace_back("avro"); |
| #endif |
| return values; |
| } |
| |
| INSTANTIATE_TEST_SUITE_P(FileFormat, AppendCompactionInteTest, |
| ::testing::ValuesIn(GetTestValuesForAppendCompactionInteTest())); |
| |
| TEST_P(AppendCompactionInteTest, TestAppendTableStreamWriteFullCompaction) { |
| auto dir = UniqueTestDirectory::Create(); |
| ASSERT_TRUE(dir); |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())}; |
| auto schema = arrow::schema(fields); |
| |
| std::vector<std::string> primary_keys = {}; |
| std::vector<std::string> partition_keys = {"f1"}; |
| auto file_format = GetParam(); |
| std::map<std::string, std::string> options = { |
| {Options::FILE_FORMAT, file_format}, |
| {Options::BUCKET, "2"}, |
| {Options::BUCKET_KEY, "f2"}, |
| {Options::FILE_SYSTEM, "local"}, |
| }; |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(dir->Str(), schema, partition_keys, primary_keys, options, |
| /*is_streaming_mode=*/true)); |
| ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema, |
| helper->LatestSchema()); |
| ASSERT_TRUE(table_schema); |
| auto gen = std::make_shared<DataGenerator>(table_schema.value(), pool_); |
| int64_t commit_identifier = 0; |
| PrepareSimpleAppendData(gen, /*with_dv=*/false, helper.get(), &commit_identifier); |
| std::vector<BinaryRow> datas_4; |
| datas_4.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Lily"), 10, 0, 17.1}, pool_.get())); |
| ASSERT_OK_AND_ASSIGN(auto batches_4, gen->SplitArrayByPartitionAndBucket(datas_4)); |
| ASSERT_EQ(1, batches_4.size()); |
| |
| ASSERT_OK(helper->write_->Write(std::move(batches_4[0]))); |
| ASSERT_OK(helper->write_->Compact(/*partition=*/{{"f1", "10"}}, /*bucket=*/1, |
| /*full_compaction=*/true)); |
| ASSERT_OK_AND_ASSIGN( |
| std::vector<std::shared_ptr<CommitMessage>> commit_messages, |
| helper->write_->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); |
| ASSERT_OK(helper->commit_->Commit(commit_messages, commit_identifier)); |
| ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot5, helper->LatestSnapshot()); |
| ASSERT_EQ(5, snapshot5.value().Id()); |
| ASSERT_EQ(11, snapshot5.value().TotalRecordCount().value()); |
| ASSERT_EQ(0, snapshot5.value().DeltaRecordCount().value()); |
| ASSERT_EQ(Snapshot::CommitKind::Compact(), snapshot5.value().GetCommitKind()); |
| ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> data_splits, |
| helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| ASSERT_EQ(data_splits.size(), 3); |
| std::map<std::pair<std::string, int32_t>, std::string> expected_datas; |
| expected_datas[std::make_pair("f1=10/", 0)] = R"([ |
| [0, "Alice", 10, 1, 11.1] |
| ])"; |
| |
| expected_datas[std::make_pair("f1=10/", 1)] = R"([ |
| [0, "Bob", 10, 0, 12.1], |
| [0, "Emily", 10, 0, 13.1], |
| [0, "Tony", 10, 0, 14.1], |
| [0, "Emily", 10, 0, 15.1], |
| [0, "Bob", 10, 0, 12.1], |
| [0, "Alex", 10, 0, 16.1], |
| [0, "David", 10, 0, 17.1], |
| [0, "Lily", 10, 0, 17.1] |
| ])"; |
| |
| expected_datas[std::make_pair("f1=20/", 0)] = R"([ |
| [0, "Lucy", 20, 1, 14.1], |
| [0, "Paul", 20, 1, null] |
| ])"; |
| |
| arrow::FieldVector fields_with_row_kind = fields; |
| fields_with_row_kind.insert(fields_with_row_kind.begin(), |
| arrow::field("_VALUE_KIND", arrow::int8())); |
| auto data_type = arrow::struct_(fields_with_row_kind); |
| |
| for (const auto& split : data_splits) { |
| auto split_impl = dynamic_cast<DataSplitImpl*>(split.get()); |
| ASSERT_OK_AND_ASSIGN(std::string partition_str, |
| helper->PartitionStr(split_impl->Partition())); |
| auto iter = expected_datas.find(std::make_pair(partition_str, split_impl->Bucket())); |
| ASSERT_TRUE(iter != expected_datas.end()); |
| ASSERT_OK_AND_ASSIGN(bool success, |
| helper->ReadAndCheckResult(data_type, {split}, iter->second)); |
| ASSERT_TRUE(success); |
| } |
| } |
| |
| TEST_P(AppendCompactionInteTest, TestAppendTableStreamWriteFullCompactionWithDv) { |
| auto dir = UniqueTestDirectory::Create(); |
| ASSERT_TRUE(dir); |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())}; |
| auto schema = arrow::schema(fields); |
| |
| std::vector<std::string> primary_keys = {}; |
| std::vector<std::string> partition_keys = {"f1"}; |
| auto file_format = GetParam(); |
| std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format}, |
| {Options::BUCKET, "2"}, |
| {Options::BUCKET_KEY, "f2"}, |
| {Options::FILE_SYSTEM, "local"}, |
| {Options::DELETION_VECTORS_ENABLED, "true"}}; |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(dir->Str(), schema, partition_keys, primary_keys, options, |
| /*is_streaming_mode=*/true)); |
| ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema, |
| helper->LatestSchema()); |
| ASSERT_TRUE(table_schema); |
| auto gen = std::make_shared<DataGenerator>(table_schema.value(), pool_); |
| int64_t commit_identifier = 0; |
| PrepareSimpleAppendData(gen, /*with_dv=*/true, helper.get(), &commit_identifier); |
| std::vector<BinaryRow> datas_4; |
| datas_4.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Lily"), 10, 0, 17.1}, pool_.get())); |
| ASSERT_OK_AND_ASSIGN(auto batches_4, gen->SplitArrayByPartitionAndBucket(datas_4)); |
| ASSERT_EQ(1, batches_4.size()); |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper2, TestHelper::Create(dir->Str(), schema, partition_keys, primary_keys, options, |
| /*is_streaming_mode=*/true, /*ignore_if_exists=*/true)); |
| |
| ASSERT_OK(helper2->write_->Write(std::move(batches_4[0]))); |
| ASSERT_OK(helper2->write_->Compact(/*partition=*/{{"f1", "10"}}, /*bucket=*/1, |
| /*full_compaction=*/true)); |
| ASSERT_OK_AND_ASSIGN( |
| std::vector<std::shared_ptr<CommitMessage>> commit_messages, |
| helper2->write_->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); |
| ASSERT_OK(helper2->commit_->Commit(commit_messages, commit_identifier)); |
| ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot5, helper2->LatestSnapshot()); |
| ASSERT_EQ(6, snapshot5.value().Id()); |
| ASSERT_EQ(8, snapshot5.value().TotalRecordCount().value()); |
| ASSERT_EQ(-3, snapshot5.value().DeltaRecordCount().value()); |
| ASSERT_EQ(Snapshot::CommitKind::Compact(), snapshot5.value().GetCommitKind()); |
| ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> data_splits, |
| helper2->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| ASSERT_EQ(data_splits.size(), 3); |
| std::map<std::pair<std::string, int32_t>, std::string> expected_datas; |
| expected_datas[std::make_pair("f1=10/", 0)] = R"([ |
| [0, "Alice", 10, 1, 11.1] |
| ])"; |
| |
| expected_datas[std::make_pair("f1=10/", 1)] = R"([ |
| [0, "Emily", 10, 0, 13.1], |
| [0, "Tony", 10, 0, 14.1], |
| [0, "Bob", 10, 0, 12.1], |
| [0, "Alex", 10, 0, 16.1], |
| [0, "Lily", 10, 0, 17.1] |
| ])"; |
| |
| expected_datas[std::make_pair("f1=20/", 0)] = R"([ |
| [0, "Lucy", 20, 1, 14.1], |
| [0, "Paul", 20, 1, null] |
| ])"; |
| |
| arrow::FieldVector fields_with_row_kind = fields; |
| fields_with_row_kind.insert(fields_with_row_kind.begin(), |
| arrow::field("_VALUE_KIND", arrow::int8())); |
| auto data_type = arrow::struct_(fields_with_row_kind); |
| |
| for (const auto& split : data_splits) { |
| auto split_impl = dynamic_cast<DataSplitImpl*>(split.get()); |
| ASSERT_OK_AND_ASSIGN(std::string partition_str, |
| helper2->PartitionStr(split_impl->Partition())); |
| auto iter = expected_datas.find(std::make_pair(partition_str, split_impl->Bucket())); |
| ASSERT_TRUE(iter != expected_datas.end()); |
| ASSERT_OK_AND_ASSIGN(bool success, |
| helper2->ReadAndCheckResult(data_type, {split}, iter->second)); |
| ASSERT_TRUE(success); |
| } |
| } |
| |
| TEST_P(AppendCompactionInteTest, TestAppendTableStreamWriteBestEffortCompaction) { |
| auto dir = UniqueTestDirectory::Create(); |
| ASSERT_TRUE(dir); |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())}; |
| auto schema = arrow::schema(fields); |
| |
| std::vector<std::string> primary_keys = {}; |
| std::vector<std::string> partition_keys = {"f1"}; |
| auto file_format = GetParam(); |
| std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format}, |
| {Options::BUCKET, "2"}, |
| {Options::BUCKET_KEY, "f2"}, |
| {Options::FILE_SYSTEM, "local"}, |
| {Options::COMPACTION_MIN_FILE_NUM, "3"}, |
| {Options::DELETION_VECTORS_ENABLED, "true"}}; |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(dir->Str(), schema, partition_keys, primary_keys, options, |
| /*is_streaming_mode=*/true)); |
| ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema, |
| helper->LatestSchema()); |
| ASSERT_TRUE(table_schema); |
| auto gen = std::make_shared<DataGenerator>(table_schema.value(), pool_); |
| int64_t commit_identifier = 0; |
| PrepareSimpleAppendData(gen, /*with_dv=*/false, helper.get(), &commit_identifier); |
| std::vector<BinaryRow> datas_4; |
| datas_4.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Lily"), 10, 0, 17.1}, pool_.get())); |
| ASSERT_OK_AND_ASSIGN(auto batches_4, gen->SplitArrayByPartitionAndBucket(datas_4)); |
| ASSERT_EQ(1, batches_4.size()); |
| |
| ASSERT_OK(helper->write_->Write(std::move(batches_4[0]))); |
| ASSERT_OK(helper->write_->Compact(/*partition=*/{{"f1", "10"}}, /*bucket=*/1, |
| /*full_compaction=*/false)); |
| ASSERT_OK_AND_ASSIGN( |
| std::vector<std::shared_ptr<CommitMessage>> commit_messages, |
| helper->write_->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); |
| ASSERT_OK(helper->commit_->Commit(commit_messages, commit_identifier)); |
| ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot5, helper->LatestSnapshot()); |
| ASSERT_EQ(5, snapshot5.value().Id()); |
| ASSERT_EQ(11, snapshot5.value().TotalRecordCount().value()); |
| ASSERT_EQ(0, snapshot5.value().DeltaRecordCount().value()); |
| ASSERT_EQ(Snapshot::CommitKind::Compact(), snapshot5.value().GetCommitKind()); |
| ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> data_splits, |
| helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| ASSERT_EQ(data_splits.size(), 3); |
| std::map<std::pair<std::string, int32_t>, std::string> expected_datas; |
| expected_datas[std::make_pair("f1=10/", 0)] = R"([ |
| [0, "Alice", 10, 1, 11.1] |
| ])"; |
| |
| expected_datas[std::make_pair("f1=10/", 1)] = R"([ |
| [0, "Bob", 10, 0, 12.1], |
| [0, "Emily", 10, 0, 13.1], |
| [0, "Tony", 10, 0, 14.1], |
| [0, "Emily", 10, 0, 15.1], |
| [0, "Bob", 10, 0, 12.1], |
| [0, "Alex", 10, 0, 16.1], |
| [0, "David", 10, 0, 17.1], |
| [0, "Lily", 10, 0, 17.1] |
| ])"; |
| |
| expected_datas[std::make_pair("f1=20/", 0)] = R"([ |
| [0, "Lucy", 20, 1, 14.1], |
| [0, "Paul", 20, 1, null] |
| ])"; |
| |
| arrow::FieldVector fields_with_row_kind = fields; |
| fields_with_row_kind.insert(fields_with_row_kind.begin(), |
| arrow::field("_VALUE_KIND", arrow::int8())); |
| auto data_type = arrow::struct_(fields_with_row_kind); |
| |
| for (const auto& split : data_splits) { |
| auto split_impl = dynamic_cast<DataSplitImpl*>(split.get()); |
| ASSERT_OK_AND_ASSIGN(std::string partition_str, |
| helper->PartitionStr(split_impl->Partition())); |
| auto iter = expected_datas.find(std::make_pair(partition_str, split_impl->Bucket())); |
| ASSERT_TRUE(iter != expected_datas.end()); |
| ASSERT_OK_AND_ASSIGN(bool success, |
| helper->ReadAndCheckResult(data_type, {split}, iter->second)); |
| ASSERT_TRUE(success); |
| } |
| } |
| |
| TEST_P(AppendCompactionInteTest, TestAppendTableStreamWriteCompactionWithExternalPath) { |
| auto dir = UniqueTestDirectory::Create(); |
| ASSERT_TRUE(dir); |
| auto external_dir = UniqueTestDirectory::Create(); |
| ASSERT_TRUE(external_dir); |
| std::string external_test_dir = "FILE://" + external_dir->Str(); |
| |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())}; |
| auto schema = arrow::schema(fields); |
| |
| std::vector<std::string> primary_keys = {}; |
| std::vector<std::string> partition_keys = {"f1"}; |
| auto file_format = GetParam(); |
| std::map<std::string, std::string> options = { |
| {Options::FILE_FORMAT, file_format}, |
| {Options::BUCKET, "2"}, |
| {Options::BUCKET_KEY, "f2"}, |
| {Options::FILE_SYSTEM, "local"}, |
| {Options::DELETION_VECTORS_ENABLED, "true"}, |
| {Options::DATA_FILE_EXTERNAL_PATHS, external_test_dir}, |
| {Options::DATA_FILE_EXTERNAL_PATHS_STRATEGY, "round-robin"}}; |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(dir->Str(), schema, partition_keys, primary_keys, options, |
| /*is_streaming_mode=*/true)); |
| ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema, |
| helper->LatestSchema()); |
| ASSERT_TRUE(table_schema); |
| auto gen = std::make_shared<DataGenerator>(table_schema.value(), pool_); |
| int64_t commit_identifier = 0; |
| PrepareSimpleAppendData(gen, /*with_dv=*/false, helper.get(), &commit_identifier); |
| std::vector<BinaryRow> datas_4; |
| datas_4.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Lily"), 10, 0, 17.1}, pool_.get())); |
| ASSERT_OK_AND_ASSIGN(auto batches_4, gen->SplitArrayByPartitionAndBucket(datas_4)); |
| ASSERT_EQ(1, batches_4.size()); |
| |
| ASSERT_OK(helper->write_->Write(std::move(batches_4[0]))); |
| ASSERT_OK(helper->write_->Compact(/*partition=*/{{"f1", "10"}}, /*bucket=*/1, |
| /*full_compaction=*/true)); |
| ASSERT_OK_AND_ASSIGN( |
| std::vector<std::shared_ptr<CommitMessage>> commit_messages, |
| helper->write_->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); |
| ASSERT_OK(helper->commit_->Commit(commit_messages, commit_identifier)); |
| ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot5, helper->LatestSnapshot()); |
| ASSERT_EQ(5, snapshot5.value().Id()); |
| ASSERT_EQ(11, snapshot5.value().TotalRecordCount().value()); |
| ASSERT_EQ(0, snapshot5.value().DeltaRecordCount().value()); |
| ASSERT_EQ(Snapshot::CommitKind::Compact(), snapshot5.value().GetCommitKind()); |
| ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> data_splits, |
| helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| ASSERT_EQ(data_splits.size(), 3); |
| std::map<std::pair<std::string, int32_t>, std::string> expected_datas; |
| expected_datas[std::make_pair("f1=10/", 0)] = R"([ |
| [0, "Alice", 10, 1, 11.1] |
| ])"; |
| |
| expected_datas[std::make_pair("f1=10/", 1)] = R"([ |
| [0, "Bob", 10, 0, 12.1], |
| [0, "Emily", 10, 0, 13.1], |
| [0, "Tony", 10, 0, 14.1], |
| [0, "Emily", 10, 0, 15.1], |
| [0, "Bob", 10, 0, 12.1], |
| [0, "Alex", 10, 0, 16.1], |
| [0, "David", 10, 0, 17.1], |
| [0, "Lily", 10, 0, 17.1] |
| ])"; |
| |
| expected_datas[std::make_pair("f1=20/", 0)] = R"([ |
| [0, "Lucy", 20, 1, 14.1], |
| [0, "Paul", 20, 1, null] |
| ])"; |
| |
| arrow::FieldVector fields_with_row_kind = fields; |
| fields_with_row_kind.insert(fields_with_row_kind.begin(), |
| arrow::field("_VALUE_KIND", arrow::int8())); |
| auto data_type = arrow::struct_(fields_with_row_kind); |
| |
| for (const auto& split : data_splits) { |
| auto split_impl = dynamic_cast<DataSplitImpl*>(split.get()); |
| ASSERT_OK_AND_ASSIGN(std::string partition_str, |
| helper->PartitionStr(split_impl->Partition())); |
| auto iter = expected_datas.find(std::make_pair(partition_str, split_impl->Bucket())); |
| ASSERT_TRUE(iter != expected_datas.end()); |
| ASSERT_OK_AND_ASSIGN(bool success, |
| helper->ReadAndCheckResult(data_type, {split}, iter->second)); |
| ASSERT_TRUE(success); |
| } |
| } |
| |
| TEST_F(AppendCompactionInteTest, TestAppendTableCompactionWithIOException) { |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())}; |
| auto schema = arrow::schema(fields); |
| |
| std::vector<std::string> primary_keys = {}; |
| std::vector<std::string> partition_keys = {"f1"}; |
| std::map<std::string, std::string> options = {{Options::FILE_FORMAT, "parquet"}, |
| {Options::BUCKET, "2"}, |
| {Options::BUCKET_KEY, "f2"}, |
| {Options::FILE_SYSTEM, "local"}, |
| {Options::DELETION_VECTORS_ENABLED, "true"}}; |
| |
| bool compaction_run_complete = false; |
| auto io_hook = IOHook::GetInstance(); |
| for (size_t i = 0; i < 600; ++i) { |
| auto dir = UniqueTestDirectory::Create(); |
| ASSERT_TRUE(dir); |
| |
| ASSERT_OK_AND_ASSIGN(auto helper, |
| TestHelper::Create(dir->Str(), schema, partition_keys, primary_keys, |
| options, /*is_streaming_mode=*/true)); |
| ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema, |
| helper->LatestSchema()); |
| ASSERT_TRUE(table_schema); |
| |
| auto gen = std::make_shared<DataGenerator>(table_schema.value(), pool_); |
| int64_t commit_identifier = 0; |
| PrepareSimpleAppendData(gen, /*with_dv=*/true, helper.get(), &commit_identifier); |
| |
| std::vector<BinaryRow> data; |
| data.push_back( |
| BinaryRowGenerator::GenerateRow({std::string("Lily"), 10, 0, 17.1}, pool_.get())); |
| ASSERT_OK_AND_ASSIGN(auto batches, gen->SplitArrayByPartitionAndBucket(data)); |
| ASSERT_EQ(1, batches.size()); |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper2, |
| TestHelper::Create(dir->Str(), schema, partition_keys, primary_keys, options, |
| /*is_streaming_mode=*/true, /*ignore_if_exists=*/true)); |
| |
| ScopeGuard guard([&io_hook]() { io_hook->Clear(); }); |
| io_hook->Reset(i, IOHook::Mode::RETURN_ERROR); |
| |
| CHECK_HOOK_STATUS(helper2->write_->Write(std::move(batches[0])), i); |
| CHECK_HOOK_STATUS(helper2->write_->Compact(/*partition=*/{{"f1", "10"}}, /*bucket=*/1, |
| /*full_compaction=*/true), |
| i); |
| |
| Result<std::vector<std::shared_ptr<CommitMessage>>> commit_messages = |
| helper2->write_->PrepareCommit(/*wait_compaction=*/true, commit_identifier); |
| CHECK_HOOK_STATUS(commit_messages.status(), i); |
| CHECK_HOOK_STATUS(helper2->commit_->Commit(commit_messages.value(), commit_identifier), i); |
| |
| compaction_run_complete = true; |
| io_hook->Clear(); |
| |
| ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> latest_snapshot, helper2->LatestSnapshot()); |
| ASSERT_TRUE(latest_snapshot); |
| ASSERT_EQ(Snapshot::CommitKind::Compact(), latest_snapshot->GetCommitKind()); |
| break; |
| } |
| |
| ASSERT_TRUE(compaction_run_complete); |
| } |
| |
| } // namespace paimon::test |
