diff --git a/src/paimon/common/global_index/union_global_index_reader.cpp b/src/paimon/common/global_index/union_global_index_reader.cpp index 0192ef167..b4894f3b4 100644 --- a/src/paimon/common/global_index/union_global_index_reader.cpp +++ b/src/paimon/common/global_index/union_global_index_reader.cpp @@ -39,84 +39,84 @@ Result> UnionGlobalIndexReader::VisitIsNull() Result> UnionGlobalIndexReader::VisitEqual( const Literal& literal) { - return Union([&literal](const std::shared_ptr& reader) { + return Union([literal](const std::shared_ptr& reader) { return reader->VisitEqual(literal); }); } Result> UnionGlobalIndexReader::VisitNotEqual( const Literal& literal) { - return Union([&literal](const std::shared_ptr& reader) { + return Union([literal](const std::shared_ptr& reader) { return reader->VisitNotEqual(literal); }); } Result> UnionGlobalIndexReader::VisitLessThan( const Literal& literal) { - return Union([&literal](const std::shared_ptr& reader) { + return Union([literal](const std::shared_ptr& reader) { return reader->VisitLessThan(literal); }); } Result> UnionGlobalIndexReader::VisitLessOrEqual( const Literal& literal) { - return Union([&literal](const std::shared_ptr& reader) { + return Union([literal](const std::shared_ptr& reader) { return reader->VisitLessOrEqual(literal); }); } Result> UnionGlobalIndexReader::VisitGreaterThan( const Literal& literal) { - return Union([&literal](const std::shared_ptr& reader) { + return Union([literal](const std::shared_ptr& reader) { return reader->VisitGreaterThan(literal); }); } Result> UnionGlobalIndexReader::VisitGreaterOrEqual( const Literal& literal) { - return Union([&literal](const std::shared_ptr& reader) { + return Union([literal](const std::shared_ptr& reader) { return reader->VisitGreaterOrEqual(literal); }); } Result> UnionGlobalIndexReader::VisitIn( const std::vector& literals) { - return Union([&literals](const std::shared_ptr& reader) { + return Union([literals](const std::shared_ptr& reader) { return reader->VisitIn(literals); }); } Result> UnionGlobalIndexReader::VisitNotIn( const std::vector& literals) { - return Union([&literals](const std::shared_ptr& reader) { + return Union([literals](const std::shared_ptr& reader) { return reader->VisitNotIn(literals); }); } Result> UnionGlobalIndexReader::VisitStartsWith( const Literal& prefix) { - return Union([&prefix](const std::shared_ptr& reader) { + return Union([prefix](const std::shared_ptr& reader) { return reader->VisitStartsWith(prefix); }); } Result> UnionGlobalIndexReader::VisitEndsWith( const Literal& suffix) { - return Union([&suffix](const std::shared_ptr& reader) { + return Union([suffix](const std::shared_ptr& reader) { return reader->VisitEndsWith(suffix); }); } Result> UnionGlobalIndexReader::VisitContains( const Literal& literal) { - return Union([&literal](const std::shared_ptr& reader) { + return Union([literal](const std::shared_ptr& reader) { return reader->VisitContains(literal); }); } Result> UnionGlobalIndexReader::VisitLike( const Literal& literal) { - return Union([&literal](const std::shared_ptr& reader) { + return Union([literal](const std::shared_ptr& reader) { return reader->VisitLike(literal); }); } @@ -124,7 +124,7 @@ Result> UnionGlobalIndexReader::VisitLike( Result> UnionGlobalIndexReader::VisitVectorSearch( const std::shared_ptr& vector_search) { auto results = ExecuteAllReaders>>( - [&vector_search](const std::shared_ptr& reader) + [vector_search](const std::shared_ptr& reader) -> Result> { return reader->VisitVectorSearch(vector_search); }); @@ -155,15 +155,13 @@ Result> UnionGlobalIndexReader::VisitVe Result> UnionGlobalIndexReader::VisitFullTextSearch( const std::shared_ptr& full_text_search) { - return Union([&full_text_search](const std::shared_ptr& reader) { + return Union([full_text_search](const std::shared_ptr& reader) { return reader->VisitFullTextSearch(full_text_search); }); } Result> UnionGlobalIndexReader::Union(ReaderAction action) { - auto results = ExecuteAllReaders>>( - [&action](const std::shared_ptr& reader) - -> Result> { return action(reader); }); + auto results = ExecuteAllReaders>>(action); std::shared_ptr merged_result = nullptr; for (auto& result_or_status : results) { @@ -207,8 +205,7 @@ std::vector UnionGlobalIndexReader::ExecuteAllReaders( std::vector> futures; futures.reserve(readers_.size()); for (const auto& reader : readers_) { - futures.push_back( - Via(executor_.get(), [&action, reader]() -> R { return action(reader); })); + futures.push_back(Via(executor_.get(), [action, reader]() -> R { return action(reader); })); } return CollectAll(futures); } diff --git a/src/paimon/common/global_index/union_global_index_reader_test.cpp b/src/paimon/common/global_index/union_global_index_reader_test.cpp index 0f6794701..248102408 100644 --- a/src/paimon/common/global_index/union_global_index_reader_test.cpp +++ b/src/paimon/common/global_index/union_global_index_reader_test.cpp @@ -19,6 +19,8 @@ #include #include #include +#include +#include #include #include @@ -54,6 +56,11 @@ class FakeReader : public GlobalIndexReader { error_message_ = message; } + void SetThrowException(const std::string& message) { + throw_exception_ = true; + exception_message_ = message; + } + /// Sets a scored result returned by VisitVectorSearch. void SetScoredResult(const std::vector& row_ids, const std::vector& scores) { scored_row_ids_ = row_ids; @@ -161,6 +168,9 @@ class FakeReader : public GlobalIndexReader { private: Result> MakeResult() { invocation_count_++; + if (throw_exception_) { + throw std::runtime_error(exception_message_); + } if (return_error_) { return Status::Invalid(error_message_); } @@ -176,7 +186,9 @@ class FakeReader : public GlobalIndexReader { std::vector default_result_; bool return_nullptr_ = false; bool return_error_ = false; + bool throw_exception_ = false; std::string error_message_; + std::string exception_message_; std::vector scored_row_ids_; std::vector scored_scores_; bool has_scored_result_ = false; @@ -184,6 +196,40 @@ class FakeReader : public GlobalIndexReader { std::atomic invocation_count_{0}; }; +// Runs the first task immediately and defers the rest. This makes it possible to verify that +// queued tasks own their action even if collecting an earlier future throws. +class DeferAfterFirstExecutor : public Executor { + public: + void Add(std::function func) override { + if (submission_count_++ == 0) { + func(); + } else { + pending_tasks_.push(std::move(func)); + } + } + + void ShutdownNow() override { + std::queue> empty; + pending_tasks_.swap(empty); + } + + uint32_t GetThreadNum() const override { + return 1; + } + + void RunPendingTasks() { + while (!pending_tasks_.empty()) { + std::function task = std::move(pending_tasks_.front()); + pending_tasks_.pop(); + task(); + } + } + + private: + uint32_t submission_count_ = 0; + std::queue> pending_tasks_; +}; + class UnionGlobalIndexReaderTest : public ::testing::Test { public: static void CheckResult(const std::shared_ptr& result, @@ -341,6 +387,23 @@ TEST_F(UnionGlobalIndexReaderTest, TestErrorPropagationWithExecutor) { ASSERT_NOK_WITH_MSG(union_reader.VisitIsNotNull(), "Unknown error for reader2"); } +TEST_F(UnionGlobalIndexReaderTest, TestDeferredTaskOwnsActionAfterEarlierFutureThrows) { + auto throwing_reader = std::make_shared(); + auto deferred_reader = std::make_shared(); + throwing_reader->SetThrowException("reader exception"); + deferred_reader->SetDefaultResult({2}); + + std::vector> readers = {throwing_reader, deferred_reader}; + auto executor = std::make_shared(); + UnionGlobalIndexReader union_reader(std::move(readers), executor); + ASSERT_THROW( + { [[maybe_unused]] auto result = union_reader.VisitIsNotNull(); }, std::runtime_error); + ASSERT_EQ(deferred_reader->InvocationCount(), 0); + + executor->RunPendingTasks(); + ASSERT_EQ(deferred_reader->InvocationCount(), 1); +} + TEST_F(UnionGlobalIndexReaderTest, TestVisitEqualUnion) { auto reader1 = std::make_shared(); auto reader2 = std::make_shared(); diff --git a/src/paimon/core/bucket/bucket_select_converter_test.cpp b/src/paimon/core/bucket/bucket_select_converter_test.cpp index 82f9c2563..106c8717b 100644 --- a/src/paimon/core/bucket/bucket_select_converter_test.cpp +++ b/src/paimon/core/bucket/bucket_select_converter_test.cpp @@ -16,12 +16,14 @@ #include "paimon/core/bucket/bucket_select_converter.h" +#include #include #include #include "arrow/api.h" #include "gtest/gtest.h" #include "paimon/core/bucket/default_bucket_function.h" +#include "paimon/core/bucket/hive_bucket_function.h" #include "paimon/core/bucket/mod_bucket_function.h" #include "paimon/data/decimal.h" #include "paimon/data/timestamp.h" @@ -34,41 +36,66 @@ namespace paimon::test { class BucketSelectConverterTest : public ::testing::Test { - protected: + public: + void AssertDefaultBucket(FieldType field_type, const Literal& literal, + const std::shared_ptr& arrow_type, + const BinaryRowGenerator::ValueType& values, + int32_t num_buckets = 17) const { + auto predicate = PredicateBuilder::Equal(0, "key", field_type, literal); + + ASSERT_OK_AND_ASSIGN( + std::optional selected_bucket, + BucketSelectConverter::Convert(predicate, {"key"}, {arrow_type}, + BucketFunctionType::DEFAULT, num_buckets, pool_.get())); + ASSERT_TRUE(selected_bucket.has_value()); + + BinaryRow row = BinaryRowGenerator::GenerateRow(values, pool_.get()); + DefaultBucketFunction function; + ASSERT_EQ(function.Bucket(row, num_buckets), selected_bucket.value()); + } + + private: std::shared_ptr pool_ = GetDefaultPool(); }; -TEST_F(BucketSelectConverterTest, SingleIntEqualDefault) { - int32_t num_buckets = 10; - Literal lit(static_cast(42)); - auto predicate = PredicateBuilder::Equal(0, "id", FieldType::INT, lit); - - ASSERT_OK_AND_ASSIGN(auto result, BucketSelectConverter::Convert( - predicate, {"id"}, {arrow::int32()}, - BucketFunctionType::DEFAULT, num_buckets, pool_.get())); - ASSERT_TRUE(result.has_value()); +TEST_F(BucketSelectConverterTest, SingleStringEqualDefault) { + std::string value = "hello_world"; + AssertDefaultBucket(FieldType::STRING, Literal(FieldType::STRING, value.c_str(), value.size()), + arrow::utf8(), {value}, 8); +} - // Verify by computing the expected bucket manually - auto row = BinaryRowGenerator::GenerateRow({static_cast(42)}, pool_.get()); - DefaultBucketFunction func; - ASSERT_EQ(func.Bucket(row, num_buckets), result.value()); +TEST_F(BucketSelectConverterTest, PrimitiveKeyTypes) { + AssertDefaultBucket(FieldType::BOOLEAN, Literal(true), arrow::boolean(), {true}); + AssertDefaultBucket(FieldType::TINYINT, Literal(static_cast(-12)), arrow::int8(), + {static_cast(-12)}); + AssertDefaultBucket(FieldType::SMALLINT, Literal(static_cast(1234)), arrow::int16(), + {static_cast(1234)}); + AssertDefaultBucket(FieldType::INT, Literal(static_cast(42)), arrow::int32(), + {static_cast(42)}, 10); + AssertDefaultBucket(FieldType::BIGINT, Literal(static_cast(123456789L)), + arrow::int64(), {static_cast(123456789L)}, 16); + AssertDefaultBucket(FieldType::FLOAT, Literal(1.25F), arrow::float32(), {1.25F}); + AssertDefaultBucket(FieldType::DOUBLE, Literal(-123.5), arrow::float64(), {-123.5}); } -TEST_F(BucketSelectConverterTest, SingleStringEqualDefault) { - int32_t num_buckets = 8; - std::string val = "hello_world"; - Literal lit(FieldType::STRING, val.c_str(), val.size()); - auto predicate = PredicateBuilder::Equal(0, "name", FieldType::STRING, lit); +TEST_F(BucketSelectConverterTest, TimestampMillisPrecision) { + // TIMESTAMP with millisecond precision (compact storage, precision=3) + Timestamp ts = Timestamp::FromEpochMillis(1700000000000L); + AssertDefaultBucket(FieldType::TIMESTAMP, Literal(ts), arrow::timestamp(arrow::TimeUnit::MILLI), + {TimestampType(ts, 3)}, 10); +} - ASSERT_OK_AND_ASSIGN(auto result, BucketSelectConverter::Convert( - predicate, {"name"}, {arrow::utf8()}, - BucketFunctionType::DEFAULT, num_buckets, pool_.get())); - ASSERT_TRUE(result.has_value()); +TEST_F(BucketSelectConverterTest, TimestampMicrosPrecision) { + // TIMESTAMP with microsecond precision (non-compact storage, precision=6) + Timestamp ts(1700000000000L, 123456); + AssertDefaultBucket(FieldType::TIMESTAMP, Literal(ts), arrow::timestamp(arrow::TimeUnit::MICRO), + {TimestampType(ts, 6)}, 10); +} - // Verify - auto row = BinaryRowGenerator::GenerateRow({val}, pool_.get()); - DefaultBucketFunction func; - ASSERT_EQ(func.Bucket(row, num_buckets), result.value()); +TEST_F(BucketSelectConverterTest, DecimalKey) { + Decimal decimal = Decimal::FromUnscaledLong(12345L, 10, 2); + AssertDefaultBucket(FieldType::DECIMAL, Literal(decimal), arrow::decimal128(10, 2), {decimal}, + 10); } TEST_F(BucketSelectConverterTest, MultiKeyAndPredicate) { @@ -173,22 +200,6 @@ TEST_F(BucketSelectConverterTest, NullPredicateReturnsNullopt) { ASSERT_FALSE(result.has_value()); } -TEST_F(BucketSelectConverterTest, BigintKeyDefault) { - int32_t num_buckets = 16; - Literal lit(static_cast(123456789L)); - auto predicate = PredicateBuilder::Equal(0, "user_id", FieldType::BIGINT, lit); - - ASSERT_OK_AND_ASSIGN(auto result, BucketSelectConverter::Convert( - predicate, {"user_id"}, {arrow::int64()}, - BucketFunctionType::DEFAULT, num_buckets, pool_.get())); - ASSERT_TRUE(result.has_value()); - - // Verify - auto row = BinaryRowGenerator::GenerateRow({static_cast(123456789L)}, pool_.get()); - DefaultBucketFunction func; - ASSERT_EQ(func.Bucket(row, num_buckets), result.value()); -} - TEST_F(BucketSelectConverterTest, AndWithExtraPredicateStillWorks) { // AND(EQUAL(id, 42), GREATER_THAN(value, 100)) // Only id is bucket key, value is not — should still derive bucket from id @@ -209,60 +220,52 @@ TEST_F(BucketSelectConverterTest, AndWithExtraPredicateStillWorks) { ASSERT_EQ(func.Bucket(row, num_buckets), result.value()); } -TEST_F(BucketSelectConverterTest, TimestampMillisPrecision) { - // TIMESTAMP with millisecond precision (compact storage, precision=3) - int32_t num_buckets = 10; - Timestamp ts = Timestamp::FromEpochMillis(1700000000000L); - Literal lit(ts); - auto predicate = PredicateBuilder::Equal(0, "ts", FieldType::TIMESTAMP, lit); +TEST_F(BucketSelectConverterTest, HiveBucketFunctionWithDecimal) { + int32_t num_buckets = 11; + Decimal decimal = Decimal::FromUnscaledLong(12345L, 10, 2); + auto int_predicate = + PredicateBuilder::Equal(0, "id", FieldType::INT, Literal(static_cast(7))); + auto decimal_predicate = + PredicateBuilder::Equal(1, "amount", FieldType::DECIMAL, Literal(decimal)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({int_predicate, decimal_predicate})); - auto arrow_type = arrow::timestamp(arrow::TimeUnit::MILLI); - ASSERT_OK_AND_ASSIGN(auto result, BucketSelectConverter::Convert( - predicate, {"ts"}, {arrow_type}, - BucketFunctionType::DEFAULT, num_buckets, pool_.get())); - ASSERT_TRUE(result.has_value()); - - // Verify: precision=3 uses compact WriteTimestamp - auto row = BinaryRowGenerator::GenerateRow({TimestampType(ts, 3)}, pool_.get()); - DefaultBucketFunction func; - ASSERT_EQ(func.Bucket(row, num_buckets), result.value()); + ASSERT_OK_AND_ASSIGN( + std::optional selected_bucket, + BucketSelectConverter::Convert(predicate, {"id", "amount"}, + {arrow::int32(), arrow::decimal128(10, 2)}, + BucketFunctionType::HIVE, num_buckets, pool_.get())); + ASSERT_TRUE(selected_bucket.has_value()); + + BinaryRow row = + BinaryRowGenerator::GenerateRow({static_cast(7), decimal}, pool_.get()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr function, + HiveBucketFunction::Create({HiveFieldInfo(FieldType::INT), + HiveFieldInfo(FieldType::DECIMAL, 10, 2)})); + ASSERT_EQ(function->Bucket(row, num_buckets), selected_bucket.value()); } -TEST_F(BucketSelectConverterTest, TimestampMicrosPrecision) { - // TIMESTAMP with microsecond precision (non-compact storage, precision=6) - int32_t num_buckets = 10; - Timestamp ts(1700000000000L, 123456); - Literal lit(ts); - auto predicate = PredicateBuilder::Equal(0, "ts", FieldType::TIMESTAMP, lit); - - auto arrow_type = arrow::timestamp(arrow::TimeUnit::MICRO); - ASSERT_OK_AND_ASSIGN(auto result, BucketSelectConverter::Convert( - predicate, {"ts"}, {arrow_type}, - BucketFunctionType::DEFAULT, num_buckets, pool_.get())); - ASSERT_TRUE(result.has_value()); +TEST_F(BucketSelectConverterTest, UnsupportedFieldTypeReturnsError) { + auto predicate = + PredicateBuilder::Equal(0, "items", FieldType::ARRAY, Literal(static_cast(42))); - // Verify: precision=6 uses non-compact WriteTimestamp (different layout than precision=3) - auto row = BinaryRowGenerator::GenerateRow({TimestampType(ts, 6)}, pool_.get()); - DefaultBucketFunction func; - ASSERT_EQ(func.Bucket(row, num_buckets), result.value()); + Result> result = + BucketSelectConverter::Convert(predicate, {"items"}, {arrow::list(arrow::int32())}, + BucketFunctionType::DEFAULT, 5, pool_.get()); + ASSERT_NOK_WITH_MSG(result.status(), "unsupported field type"); } -TEST_F(BucketSelectConverterTest, DecimalKey) { - int32_t num_buckets = 10; - Decimal dec = Decimal::FromUnscaledLong(12345L, 10, 2); - Literal lit(dec); - auto predicate = PredicateBuilder::Equal(0, "amount", FieldType::DECIMAL, lit); - - auto arrow_type = arrow::decimal128(10, 2); - ASSERT_OK_AND_ASSIGN(auto result, BucketSelectConverter::Convert( - predicate, {"amount"}, {arrow_type}, - BucketFunctionType::DEFAULT, num_buckets, pool_.get())); - ASSERT_TRUE(result.has_value()); - - // Verify - auto row = BinaryRowGenerator::GenerateRow({dec}, pool_.get()); - DefaultBucketFunction func; - ASSERT_EQ(func.Bucket(row, num_buckets), result.value()); +TEST_F(BucketSelectConverterTest, ModBucketFunctionWithMultipleKeysReturnsError) { + auto id_predicate = + PredicateBuilder::Equal(0, "id", FieldType::INT, Literal(static_cast(42))); + auto region_predicate = + PredicateBuilder::Equal(1, "region", FieldType::INT, Literal(static_cast(1))); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({id_predicate, region_predicate})); + + Result> result = BucketSelectConverter::Convert( + predicate, {"id", "region"}, {arrow::int32(), arrow::int32()}, BucketFunctionType::MOD, 5, + pool_.get()); + ASSERT_NOK_WITH_MSG(result.status(), + "MOD bucket function requires exactly one bucket key field"); } } // namespace paimon::test diff --git a/src/paimon/core/bucket/hive_bucket_function_test.cpp b/src/paimon/core/bucket/hive_bucket_function_test.cpp index 6ac49e607..426c3df90 100644 --- a/src/paimon/core/bucket/hive_bucket_function_test.cpp +++ b/src/paimon/core/bucket/hive_bucket_function_test.cpp @@ -105,6 +105,11 @@ class HiveBucketFunctionTest : public ::testing::Test { return BinaryRowGenerator::GenerateRow({value}, pool.get()); } + BinaryRow CreateShortRow(int16_t value) { + auto pool = GetDefaultPool(); + return BinaryRowGenerator::GenerateRow({value}, pool.get()); + } + float FloatFromBits(uint32_t bits) { float value; std::memcpy(&value, &bits, sizeof(value)); @@ -256,6 +261,14 @@ TEST_F(HiveBucketFunctionTest, TestTinyintNegativeValuesCompatibleWithJava) { ASSERT_EQ(520, func->Bucket(CreateByteRow(std::numeric_limits::min()), 1000)); } +TEST_F(HiveBucketFunctionTest, TestSmallintField) { + std::vector field_types = {FieldType::SMALLINT}; + ASSERT_OK_AND_ASSIGN(auto func, HiveBucketFunction::Create(field_types)); + + ASSERT_EQ(234, func->Bucket(CreateShortRow(static_cast(1234)), 1000)); + ASSERT_EQ(647, func->Bucket(CreateShortRow(static_cast(-1)), 1000)); +} + /// Test STRING field TEST_F(HiveBucketFunctionTest, TestStringField) { std::vector field_types = {FieldType::STRING}; diff --git a/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg_test.cpp index f8d02022f..6c1a6ffd3 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg_test.cpp @@ -118,6 +118,14 @@ TEST_F(FieldListaggAggTest, TestDistinctNoDuplicates) { ASSERT_EQ(DataDefine::GetVariantValue(ret), "a b c d"); } +TEST_F(FieldListaggAggTest, TestDistinctWithEmptyDelimiterFallsBackToWhitespace) { + ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg("", true)); + + // Empty delimiter falls back to whitespace, so the repeated "b" is removed. + auto ret = agg->Agg(std::string_view("a b"), std::string_view("b c")); + ASSERT_EQ(DataDefine::GetVariantValue(ret), "a b c"); +} + TEST_F(FieldListaggAggTest, TestDistinctEmptyInput) { ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg(";", true)); diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_test.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_test.cpp index 2f9b6bbbb..531f211e6 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_test.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_test.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -28,12 +29,17 @@ #include "arrow/api.h" #include "gtest/gtest.h" #include "paimon/common/utils/fields_comparator.h" +#include "paimon/core/compact/compact_deletion_file.h" +#include "paimon/core/deletionvectors/bucketed_dv_maintainer.h" +#include "paimon/core/deletionvectors/deletion_vectors_index_file.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/mergetree/level_sorted_run.h" #include "paimon/core/mergetree/levels.h" #include "paimon/core/mergetree/sorted_run.h" #include "paimon/core/stats/simple_stats.h" +#include "paimon/fs/file_system_factory.h" +#include "paimon/testing/mock/mock_index_path_factory.h" #include "paimon/testing/utils/binary_row_generator.h" #include "paimon/testing/utils/testharness.h" @@ -409,6 +415,80 @@ TEST_F(MergeTreeCompactManagerTest, TestTriggerFullCompaction) { } } +TEST_F(MergeTreeCompactManagerTest, TestFullCompactionRewritesEachMaxLevelFile) { + std::vector> files = { + ToFile(LevelMinMax(2, 1, 3), /*max_sequence=*/3), + ToFile(LevelMinMax(2, 4, 6), /*max_sequence=*/6)}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr levels, CreateLevels(files)); + + auto manager = std::make_shared( + levels, std::make_shared(TestStrategy()), comparator_, + /*compaction_file_size=*/2, + /*num_sorted_run_stop_trigger=*/std::numeric_limits::max(), + std::make_shared(/*expected_drop_delete=*/true), + /*metrics_reporter=*/nullptr, + /*dv_maintainer=*/nullptr, + /*lazy_gen_deletion_file=*/false, + /*need_lookup=*/false, + /*force_rewrite_all_files=*/true, + /*force_keep_delete=*/false, std::make_shared(), + std::make_shared()); + + ASSERT_OK(manager->TriggerCompaction(/*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN(std::optional> compact_result, + manager->GetCompactionResult(/*blocking=*/true)); + ASSERT_TRUE(compact_result.has_value()); + ASSERT_EQ(compact_result.value()->Before(), files); + ASSERT_EQ(compact_result.value()->After().size(), 2); + for (const auto& file : compact_result.value()->After()) { + ASSERT_EQ(file->file_name.rfind("rewrite-", /*pos=*/0), 0); + } +} + +TEST_F(MergeTreeCompactManagerTest, TestCompactionGeneratesDeletionFileEagerly) { + std::vector> files = { + ToFile(LevelMinMax(0, 1, 3), /*max_sequence=*/0), + ToFile(LevelMinMax(1, 1, 5), /*max_sequence=*/1)}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr levels, CreateLevels(files)); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr fs, + FileSystemFactory::Get("local", dir->Str(), {})); + auto path_factory = std::make_shared(dir->Str()); + auto dv_index_file = + std::make_shared(fs, path_factory, /*bitmap64=*/false, pool_); + auto dv_maintainer = std::make_shared( + dv_index_file, std::map>{}); + ASSERT_OK(dv_maintainer->NotifyNewDeletion("remaining-file", /*position=*/0)); + + auto manager = std::make_shared( + levels, std::make_shared(TestStrategy()), comparator_, + /*compaction_file_size=*/2, + /*num_sorted_run_stop_trigger=*/std::numeric_limits::max(), + std::make_shared(/*expected_drop_delete=*/true), + /*metrics_reporter=*/nullptr, dv_maintainer, + /*lazy_gen_deletion_file=*/false, + /*need_lookup=*/false, + /*force_rewrite_all_files=*/false, + /*force_keep_delete=*/false, std::make_shared(), + std::make_shared()); + + ASSERT_OK(manager->TriggerCompaction(/*full_compaction=*/false)); + ASSERT_OK_AND_ASSIGN(std::optional> compact_result, + manager->GetCompactionResult(/*blocking=*/true)); + ASSERT_TRUE(compact_result.has_value()); + std::shared_ptr deletion_file = compact_result.value()->DeletionFile(); + ASSERT_NE(deletion_file, nullptr); + ASSERT_NE(std::dynamic_pointer_cast(deletion_file), nullptr); + + ASSERT_OK_AND_ASSIGN(std::optional> index_file, + deletion_file->GetOrCompute()); + ASSERT_TRUE(index_file.has_value()); + ASSERT_EQ(index_file.value()->IndexType(), DeletionVectorsIndexFile::DELETION_VECTORS_INDEX); + ASSERT_OK_AND_ASSIGN(bool exists, dv_index_file->Exists(index_file.value())); + ASSERT_TRUE(exists); +} + TEST_F(MergeTreeCompactManagerTest, TestRejectReentrantFullCompaction) { std::vector inputs = {LevelMinMax(0, 1, 3), LevelMinMax(1, 2, 5), LevelMinMax(1, 6, 7)}; diff --git a/src/paimon/core/operation/commit/commit_scanner_test.cpp b/src/paimon/core/operation/commit/commit_scanner_test.cpp index a46364002..99d788630 100644 --- a/src/paimon/core/operation/commit/commit_scanner_test.cpp +++ b/src/paimon/core/operation/commit/commit_scanner_test.cpp @@ -28,9 +28,15 @@ #include "paimon/common/data/binary_row_writer.h" #include "paimon/common/utils/binary_row_partition_computer.h" #include "paimon/core/core_options.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/manifest/file_kind.h" +#include "paimon/core/manifest/index_manifest_entry.h" +#include "paimon/core/manifest/index_manifest_file.h" #include "paimon/core/manifest/manifest_entry.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/snapshot.h" +#include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/format/file_format_factory.h" #include "paimon/scan_context.h" #include "paimon/testing/utils/testharness.h" @@ -38,7 +44,7 @@ namespace paimon::test { namespace { -Snapshot MakeSnapshot() { +Snapshot MakeSnapshot(std::optional index_manifest = std::nullopt) { return Snapshot( /*id=*/1, /*schema_id=*/0, @@ -48,7 +54,7 @@ Snapshot MakeSnapshot() { /*delta_manifest_list_size=*/std::nullopt, /*changelog_manifest_list=*/std::nullopt, /*changelog_manifest_list_size=*/std::nullopt, - /*index_manifest=*/std::nullopt, + /*index_manifest=*/std::move(index_manifest), /*commit_user=*/"test-user", /*commit_identifier=*/1, Snapshot::CommitKind::Append(), /*time_millis=*/0, @@ -69,11 +75,22 @@ BinaryRow CreateIntPartition(int32_t value) { return row; } +IndexManifestEntry CreateIndexEntry(const std::string& file_name, int32_t partition_value) { + auto index_file = std::make_shared( + /*index_type=*/"HASH", file_name, /*file_size=*/10, /*row_count=*/1, + /*dv_ranges=*/std::nullopt, + /*external_path=*/std::nullopt); + return IndexManifestEntry(FileKind::Add(), CreateIntPartition(partition_value), + /*bucket=*/0, index_file); +} + } // namespace class CommitScannerTest : public testing::Test { protected: void SetUp() override { + dir_ = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir_); schema_ = arrow::schema({arrow::field("pt", arrow::int32())}); ASSERT_OK_AND_ASSIGN(core_options_, CoreOptions::FromMap({})); ASSERT_OK_AND_ASSIGN(partition_computer_, @@ -83,19 +100,41 @@ class CommitScannerTest : public testing::Test { /*legacy_partition_name_enabled=*/true, GetDefaultPool())); } - CommitScanner CreateScanner(CommitScanner::ScanSupplier scan_supplier) const { + CommitScanner CreateScanner( + CommitScanner::ScanSupplier scan_supplier, + const std::shared_ptr& index_manifest_file = nullptr) const { return CommitScanner( /*snapshot_manager=*/nullptr, /*schema_manager=*/nullptr, /*manifest_list=*/nullptr, /*manifest_file=*/nullptr, - /*index_manifest_file=*/nullptr, + /*index_manifest_file=*/index_manifest_file, /*table_schema=*/nullptr, schema_, core_options_, /*executor=*/nullptr, GetDefaultPool(), partition_computer_.get(), std::move(scan_supplier)); } + Result> CreateIndexManifestFile() const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_format, + FileFormatFactory::Get("orc", {})); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr path_factory, + FileStorePathFactory::Create( + dir_->Str(), schema_, /*partition_keys=*/{"pt"}, + /*default_part_value=*/"__DEFAULT_PARTITION__", file_format->Identifier(), + /*data_file_prefix=*/"data-", + /*legacy_partition_name_enabled=*/true, /*external_paths=*/{}, + /*global_index_external_path=*/std::nullopt, + /*index_file_in_data_file_dir=*/false, GetDefaultPool())); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr index_manifest_file, + IndexManifestFile::Create(dir_->GetFileSystem(), file_format, "zstd", path_factory, + /*bucket_mode=*/2, GetDefaultPool(), core_options_)); + return std::shared_ptr(std::move(index_manifest_file)); + } + protected: + std::unique_ptr dir_; std::shared_ptr schema_; CoreOptions core_options_; std::unique_ptr partition_computer_; @@ -147,4 +186,30 @@ TEST_F(CommitScannerTest, TestReadAllEntriesFromChangedPartitionsBuildsScanFilte ASSERT_EQ("42", captured_partition_filters[0]["pt"]); } +TEST_F(CommitScannerTest, TestReadAllIndexEntriesFromPartitions) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr index_manifest_file, + CreateIndexManifestFile()); + std::vector entries = {CreateIndexEntry("index-1", /*partition_value=*/1), + CreateIndexEntry("index-2", /*partition_value=*/2), + CreateIndexEntry("index-3", /*partition_value=*/3)}; + ASSERT_OK_AND_ASSIGN( + std::optional index_manifest, + index_manifest_file->WriteIndexFiles(/*previous_index_manifest=*/std::nullopt, entries)); + ASSERT_TRUE(index_manifest); + + CommitScanner scanner = CreateScanner(CommitScanner::ScanSupplier{}, index_manifest_file); + Snapshot snapshot = MakeSnapshot(index_manifest); + + ASSERT_OK_AND_ASSIGN(std::vector unfiltered, + scanner.ReadAllIndexEntriesFromPartitions(snapshot, /*partitions=*/{})); + ASSERT_EQ(3u, unfiltered.size()); + + std::vector> partitions = { + {{"pt", "2"}}, {{"unknown_partition_key", "value"}}}; + ASSERT_OK_AND_ASSIGN(std::vector filtered, + scanner.ReadAllIndexEntriesFromPartitions(snapshot, partitions)); + ASSERT_EQ(1u, filtered.size()); + ASSERT_EQ("index-2", filtered[0].index_file->FileName()); +} + } // namespace paimon::test diff --git a/src/paimon/core/operation/manifest_file_merger_test.cpp b/src/paimon/core/operation/manifest_file_merger_test.cpp index b06e44d73..60b13fe67 100644 --- a/src/paimon/core/operation/manifest_file_merger_test.cpp +++ b/src/paimon/core/operation/manifest_file_merger_test.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -108,9 +109,22 @@ class ManifestFileMergerTest : public testing::Test { return manifest_file_metas[0]; } + std::set ListManifestFiles() const { + std::vector> file_statuses; + EXPECT_OK(file_system_->ListFileStatus( + FileStorePathFactory::ManifestPath(path_factory_->RootPath()), &file_statuses)); + std::set files; + for (const auto& status : file_statuses) { + if (!status->IsDir()) { + files.insert(status->GetPath()); + } + } + return files; + } + private: void CreateManifestFile(const std::string& path_str) { - auto file_system = std::make_shared(); + file_system_ = std::make_shared(); ASSERT_OK_AND_ASSIGN( std::shared_ptr file_format, FileFormatFactory::Get("parquet", std::map())); @@ -128,10 +142,11 @@ class ManifestFileMergerTest : public testing::Test { options.GetFileFormat()->Identifier(), options.DataFilePrefix(), options.LegacyPartitionNameEnabled(), external_paths, global_index_external_path, options.IndexFileInDataFileDir(), pool_)); + path_factory_ = path_factory; ASSERT_OK_AND_ASSIGN(std::shared_ptr partition_schema, FieldMapping::GetPartitionSchema(schema, {"f0"})); ASSERT_OK_AND_ASSIGN(manifest_file_, - ManifestFile::Create(file_system, file_format, "zstd", path_factory, + ManifestFile::Create(file_system_, file_format, "zstd", path_factory, /*target_file_size=*/1024 * 1024, pool_, options, partition_schema)); } @@ -165,6 +180,8 @@ class ManifestFileMergerTest : public testing::Test { std::string test_root_; std::shared_ptr pool_; std::shared_ptr manifest_file_; + std::shared_ptr file_system_; + std::shared_ptr path_factory_; std::shared_ptr partition_type_; }; @@ -369,4 +386,24 @@ TEST_F(ManifestFileMergerTest, TestTriggerFullCompaction) { ContainSameEntryFile(merged.value(), entry_file_expected); } +TEST_F(ManifestFileMergerTest, TestDeleteNewManifestFilesWhenMinorCompactionFails) { + std::vector input; + for (int32_t i = 0; i < 4; i++) { + input.push_back(MakeManifest({MakeEntry(FileKind::Add(), std::to_string(i))})); + } + + // The first pair is merged successfully. Removing a source file from the second pair makes + // the following merge fail after a new manifest file has already been created. + std::string missing_manifest_path = path_factory_->ToManifestFilePath(input[2].FileName()); + ASSERT_OK(file_system_->Delete(missing_manifest_path)); + std::set files_before_merge = ListManifestFiles(); + + ASSERT_NOK_WITH_MSG(ManifestFileMerger::Merge( + input, /*manifest_target_file_size=*/5000, /*merge_min_count=*/2, + /*full_compaction_file_size=*/MAX_LONG_VALUE, manifest_file_.get()), + "not exists"); + + ASSERT_EQ(files_before_merge, ListManifestFiles()); +} + } // namespace paimon::test diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index e794a403d..86df4841d 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -263,10 +263,11 @@ Result> MergeFileSplitRead::CreateNoMergeReader( pool_); // create read schema without extra fields (e.g., completed key, sequence fields) - auto row_kind_field = DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()); - - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, - raw_read_schema_->AddField(0, row_kind_field)); + std::shared_ptr read_schema = raw_read_schema_; + if (read_schema->GetFieldIndex(SpecialFields::ValueKind().Name()) < 0) { + auto row_kind_field = DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(read_schema, read_schema->AddField(0, row_kind_field)); + } PAIMON_ASSIGN_OR_RAISE( std::vector> raw_file_readers, CreateRawFileReaders(data_split->Partition(), data_split->DataFiles(), read_schema, diff --git a/src/paimon/core/table/system/audit_log_system_table.cpp b/src/paimon/core/table/system/audit_log_system_table.cpp index 488ec4796..6f4432c98 100644 --- a/src/paimon/core/table/system/audit_log_system_table.cpp +++ b/src/paimon/core/table/system/audit_log_system_table.cpp @@ -31,6 +31,7 @@ #include "arrow/c/bridge.h" #include "arrow/util/checked_cast.h" #include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/types/row_kind.h" @@ -38,6 +39,7 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/core/core_options.h" #include "paimon/core/schema/table_schema.h" +#include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/key_value_table_read.h" #include "paimon/defs.h" #include "paimon/read_context.h" @@ -52,7 +54,9 @@ namespace { class AuditLogBatchConverter : public ChangelogBatchConverter { public: Result> ConvertDataColumn( - const std::shared_ptr& array, arrow::MemoryPool* /*pool*/) const override { + const std::shared_ptr& array, + const std::vector& /*row_group_lengths*/, + arrow::MemoryPool* /*pool*/) const override { return array; } }; @@ -62,64 +66,76 @@ class ChangelogBatchReader : public BatchReader { ChangelogBatchReader(std::unique_ptr reader, std::shared_ptr output_schema, bool include_sequence_number, std::shared_ptr converter, - const std::shared_ptr& pool) + bool pack_update_before_after, const std::shared_ptr& pool) : reader_(std::move(reader)), output_schema_(std::move(output_schema)), include_sequence_number_(include_sequence_number), converter_(std::move(converter)), + pack_update_before_after_(pack_update_before_after), arrow_pool_holder_(GetArrowPool(pool)), arrow_pool_(arrow_pool_holder_.get()) {} Result NextBatch() override { - PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, reader_->NextBatch()); - if (BatchReader::IsEofBatch(batch)) { - return batch; - } - auto& [c_array, c_schema] = batch; - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_array, - arrow::ImportArray(c_array.get(), c_schema.get())); - std::shared_ptr struct_array = - std::dynamic_pointer_cast(arrow_array); - if (!struct_array) { - return Status::Invalid("audit_log system table expects struct batches"); - } + while (true) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, reader_->NextBatch()); + std::shared_ptr struct_array; + std::vector row_group_lengths; + if (BatchReader::IsEofBatch(batch)) { + if (!pending_update_before_) { + return batch; + } + struct_array = std::move(pending_update_before_); + row_group_lengths.push_back(1); + } else { + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr arrow_array, + arrow::ImportArray(c_array.get(), c_schema.get())); + struct_array = std::dynamic_pointer_cast(arrow_array); + if (!struct_array) { + return Status::Invalid("audit_log system table expects struct batches"); + } + PAIMON_ASSIGN_OR_RAISE(struct_array, PrependPendingUpdateBefore(struct_array)); + PAIMON_ASSIGN_OR_RAISE(row_group_lengths, BuildRowGroupLengths(struct_array)); + } + if (row_group_lengths.empty()) { + continue; + } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr rowkind_array, - BuildRowKindArray(struct_array)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr row_kind_array, + BuildRowKindArray(struct_array, row_group_lengths)); - arrow::ArrayVector output_arrays = {rowkind_array}; - if (include_sequence_number_) { - std::shared_ptr sequence_array = - struct_array->GetFieldByName(SpecialFields::SequenceNumber().Name()); - if (!sequence_array) { - return Status::Invalid("cannot find _SEQUENCE_NUMBER in audit_log batch"); + arrow::ArrayVector output_arrays = {row_kind_array}; + if (include_sequence_number_) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr sequence_array, + BuildSequenceNumberArray(struct_array, row_group_lengths)); + output_arrays.push_back(sequence_array); } - PAIMON_ASSIGN_OR_RAISE(sequence_array, CopyToStablePool(sequence_array)); - output_arrays.push_back(sequence_array); - } - for (const auto& field : output_schema_->fields()) { - if (field->name() == SpecialFields::RowKind().Name() || - field->name() == SpecialFields::SequenceNumber().Name()) { - continue; - } - std::shared_ptr array = struct_array->GetFieldByName(field->name()); - if (!array) { - return Status::Invalid("cannot find ", field->name(), " in changelog batch"); + for (const auto& field : output_schema_->fields()) { + if (field->name() == SpecialFields::RowKind().Name() || + field->name() == SpecialFields::SequenceNumber().Name()) { + continue; + } + std::shared_ptr array = struct_array->GetFieldByName(field->name()); + if (!array) { + return Status::Invalid("cannot find ", field->name(), " in changelog batch"); + } + PAIMON_ASSIGN_OR_RAISE( + array, converter_->ConvertDataColumn(array, row_group_lengths, arrow_pool_)); + PAIMON_ASSIGN_OR_RAISE(array, CopyToStablePool(array)); + output_arrays.push_back(array); } - PAIMON_ASSIGN_OR_RAISE(array, converter_->ConvertDataColumn(array, arrow_pool_)); - PAIMON_ASSIGN_OR_RAISE(array, CopyToStablePool(array)); - output_arrays.push_back(array); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr output_array, - arrow::StructArray::Make(output_arrays, output_schema_->field_names())); - auto output_c_array = std::make_unique(); - auto output_c_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW( - arrow::ExportArray(*output_array, output_c_array.get(), output_c_schema.get())); - return std::make_pair(std::move(output_c_array), std::move(output_c_schema)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr output_array, + arrow::StructArray::Make(output_arrays, output_schema_->field_names())); + auto output_c_array = std::make_unique(); + auto output_c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*output_array, output_c_array.get(), output_c_schema.get())); + return std::make_pair(std::move(output_c_array), std::move(output_c_schema)); + } } std::shared_ptr GetReaderMetrics() const override { @@ -127,10 +143,68 @@ class ChangelogBatchReader : public BatchReader { } void Close() override { + pending_update_before_.reset(); reader_->Close(); } private: + Result> PrependPendingUpdateBefore( + const std::shared_ptr& struct_array) { + if (!pending_update_before_) { + return struct_array; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr combined, + arrow::Concatenate({pending_update_before_, struct_array}, arrow_pool_)); + pending_update_before_.reset(); + std::shared_ptr result = + std::dynamic_pointer_cast(combined); + if (!result) { + return Status::Invalid("failed to concatenate binlog struct batches"); + } + return result; + } + + Result> BuildRowGroupLengths( + const std::shared_ptr& struct_array) { + std::shared_ptr value_kind_array = + std::dynamic_pointer_cast( + struct_array->GetFieldByName(SpecialFields::ValueKind().Name())); + if (!value_kind_array) { + return Status::Invalid("cannot find _VALUE_KIND in audit_log batch"); + } + + std::vector row_group_lengths; + row_group_lengths.reserve(struct_array->length()); + for (int64_t i = 0; i < value_kind_array->length();) { + bool is_update_before = + !value_kind_array->IsNull(i) && + value_kind_array->Value(i) == RowKind::UpdateBefore()->ToByteValue(); + if (!pack_update_before_after_ || !is_update_before) { + row_group_lengths.push_back(1); + ++i; + continue; + } + if (i + 1 == value_kind_array->length()) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr pending, + CopyToStablePool(struct_array->Slice(i, /*length=*/1))); + pending_update_before_ = std::dynamic_pointer_cast(pending); + if (!pending_update_before_) { + return Status::Invalid("failed to cache UPDATE_BEFORE in binlog reader"); + } + break; + } + if (value_kind_array->IsNull(i + 1) || + value_kind_array->Value(i + 1) != RowKind::UpdateAfter()->ToByteValue()) { + return Status::Invalid( + "UPDATE_BEFORE is not followed by UPDATE_AFTER in binlog reader"); + } + row_group_lengths.push_back(2); + i += 2; + } + return row_group_lengths; + } + Result> CopyToStablePool( const std::shared_ptr& array) const { /// The imported data batch may release its C Arrow buffers after this wrapper returns. @@ -141,7 +215,8 @@ class ChangelogBatchReader : public BatchReader { } Result> BuildRowKindArray( - const std::shared_ptr& struct_array) const { + const std::shared_ptr& struct_array, + const std::vector& row_group_lengths) const { std::shared_ptr value_kind_array = std::dynamic_pointer_cast( struct_array->GetFieldByName(SpecialFields::ValueKind().Name())); @@ -149,15 +224,46 @@ class ChangelogBatchReader : public BatchReader { return Status::Invalid("cannot find _VALUE_KIND in audit_log batch"); } arrow::StringBuilder builder(arrow_pool_); - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(value_kind_array->length())); - for (int64_t i = 0; i < value_kind_array->length(); ++i) { - if (value_kind_array->IsNull(i)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(row_group_lengths.size())); + int64_t offset = 0; + for (int32_t row_group_length : row_group_lengths) { + int64_t row_kind_index = offset + row_group_length - 1; + if (value_kind_array->IsNull(row_kind_index)) { + return Status::Invalid( + fmt::format("exists null value in value kind array in pos {}", row_kind_index)); PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); - continue; + } else { + PAIMON_ASSIGN_OR_RAISE( + const RowKind* row_kind, + RowKind::FromByteValue(value_kind_array->Value(row_kind_index))); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(row_kind->ShortString())); + } + offset += row_group_length; + } + std::shared_ptr result; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&result)); + return result; + } + + Result> BuildSequenceNumberArray( + const std::shared_ptr& struct_array, + const std::vector& row_group_lengths) const { + std::shared_ptr sequence_array = + std::dynamic_pointer_cast( + struct_array->GetFieldByName(SpecialFields::SequenceNumber().Name())); + if (!sequence_array) { + return Status::Invalid("cannot find _SEQUENCE_NUMBER in audit_log batch"); + } + arrow::Int64Builder builder(arrow_pool_); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(row_group_lengths.size())); + int64_t offset = 0; + for (int32_t row_group_length : row_group_lengths) { + if (sequence_array->IsNull(offset)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(sequence_array->Value(offset))); } - PAIMON_ASSIGN_OR_RAISE(const RowKind* row_kind, - RowKind::FromByteValue(value_kind_array->Value(i))); - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(row_kind->ShortString())); + offset += row_group_length; } std::shared_ptr result; PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&result)); @@ -168,8 +274,10 @@ class ChangelogBatchReader : public BatchReader { std::shared_ptr output_schema_; bool include_sequence_number_; std::shared_ptr converter_; + bool pack_update_before_after_; std::unique_ptr arrow_pool_holder_; arrow::MemoryPool* arrow_pool_; + std::shared_ptr pending_update_before_; }; class ChangelogTableRead : public TableRead { @@ -186,17 +294,42 @@ class ChangelogTableRead : public TableRead { Result> CreateReader( const std::vector>& splits) override { + // Records across different splits should not be packed, because for streaming reads on a + // primary-key table, all data belonging to the same partition and bucket is placed in a + // single split. Therefore, an UPDATE_BEFORE/UPDATE_AFTER pair will not be truncated at a + // split boundary. + if (converter_->PackUpdateBeforeAfter()) { + std::vector> readers; + readers.reserve(splits.size()); + for (const auto& split : splits) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, CreateReader(split)); + readers.push_back(std::move(reader)); + } + return std::make_unique(std::move(readers), GetMemoryPool()); + } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, data_read_->CreateReader(splits)); - return std::make_unique(std::move(reader), output_schema_, - include_sequence_number_, converter_, - GetMemoryPool()); + return CreateChangelogBatchReader(std::move(reader), output_schema_, + include_sequence_number_, converter_, + /*pack_update_before_after=*/false, GetMemoryPool()); } Result> CreateReader( const std::shared_ptr& split) override { - std::vector> splits = {split}; - return CreateReader(splits); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + data_read_->CreateReader(split)); + bool pack_update_before_after = false; + if (converter_->PackUpdateBeforeAfter()) { + std::shared_ptr data_split = + std::dynamic_pointer_cast(split); + if (!data_split) { + return Status::Invalid("binlog system table expects data split"); + } + pack_update_before_after = data_split->IsStreaming(); + } + return CreateChangelogBatchReader(std::move(reader), output_schema_, + include_sequence_number_, converter_, + pack_update_before_after, GetMemoryPool()); } private: @@ -208,6 +341,15 @@ class ChangelogTableRead : public TableRead { } // namespace +std::unique_ptr CreateChangelogBatchReader( + std::unique_ptr reader, std::shared_ptr output_schema, + bool include_sequence_number, std::shared_ptr converter, + bool pack_update_before_after, const std::shared_ptr& pool) { + return std::make_unique(std::move(reader), std::move(output_schema), + include_sequence_number, std::move(converter), + pack_update_before_after, pool); +} + AuditLogSystemTable::AuditLogSystemTable(std::shared_ptr fs, std::string table_path, std::shared_ptr table_schema, std::map options) @@ -221,10 +363,10 @@ std::string AuditLogSystemTable::Name() const { } Result> AuditLogSystemTable::ArrowSchema() const { - std::shared_ptr rowkind_field = + std::shared_ptr row_kind_field = DataField::ConvertDataFieldToArrowField(SpecialFields::RowKind()); - rowkind_field = rowkind_field->WithNullable(false); - arrow::FieldVector fields = {rowkind_field}; + row_kind_field = row_kind_field->WithNullable(false); + arrow::FieldVector fields = {row_kind_field}; PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options_)); if (core_options.TableReadSequenceNumberEnabled()) { fields.push_back(DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())); diff --git a/src/paimon/core/table/system/audit_log_system_table.h b/src/paimon/core/table/system/audit_log_system_table.h index ef7baa983..354968dc0 100644 --- a/src/paimon/core/table/system/audit_log_system_table.h +++ b/src/paimon/core/table/system/audit_log_system_table.h @@ -33,9 +33,19 @@ class ChangelogBatchConverter { virtual ~ChangelogBatchConverter() = default; virtual Result> ConvertDataColumn( - const std::shared_ptr& array, arrow::MemoryPool* pool) const = 0; + const std::shared_ptr& array, const std::vector& row_group_lengths, + arrow::MemoryPool* pool) const = 0; + + virtual bool PackUpdateBeforeAfter() const { + return false; + } }; +std::unique_ptr CreateChangelogBatchReader( + std::unique_ptr reader, std::shared_ptr output_schema, + bool include_sequence_number, std::shared_ptr converter, + bool pack_update_before_after, const std::shared_ptr& pool); + /// System table for `T$audit_log`, exposing row-level changelog records with rowkind. class AuditLogSystemTable : public SystemTable { public: diff --git a/src/paimon/core/table/system/binlog_system_table.cpp b/src/paimon/core/table/system/binlog_system_table.cpp index 7dcb58d68..76e8f9090 100644 --- a/src/paimon/core/table/system/binlog_system_table.cpp +++ b/src/paimon/core/table/system/binlog_system_table.cpp @@ -38,11 +38,15 @@ namespace { class BinlogBatchConverter : public ChangelogBatchConverter { public: Result> ConvertDataColumn( - const std::shared_ptr& array, arrow::MemoryPool* pool) const override { + const std::shared_ptr& array, const std::vector& row_group_lengths, + arrow::MemoryPool* pool) const override { arrow::Int32Builder offsets_builder(pool); - PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets_builder.Reserve(array->length() + 1)); - for (int64_t i = 0; i <= array->length(); ++i) { - PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets_builder.Append(static_cast(i))); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets_builder.Reserve(row_group_lengths.size() + 1)); + int32_t offset = 0; + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets_builder.Append(offset)); + for (int32_t row_group_length : row_group_lengths) { + offset += row_group_length; + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets_builder.Append(offset)); } std::shared_ptr offsets_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets_builder.Finish(&offsets_array)); @@ -51,10 +55,18 @@ class BinlogBatchConverter : public ChangelogBatchConverter { arrow::ListArray::FromArrays(*offsets_array, *array, pool)); return list_array; } + + bool PackUpdateBeforeAfter() const override { + return true; + } }; } // namespace +std::shared_ptr CreateBinlogBatchConverter() { + return std::make_shared(); +} + BinlogSystemTable::BinlogSystemTable(std::shared_ptr fs, std::string table_path, std::shared_ptr table_schema, std::map options) @@ -84,7 +96,7 @@ Result> BinlogSystemTable::ArrowSchema() const { Result> BinlogSystemTable::NewRead( const std::shared_ptr& context) const { - return NewChangelogRead(context, std::make_shared()); + return NewChangelogRead(context, CreateBinlogBatchConverter()); } } // namespace paimon diff --git a/src/paimon/core/table/system/binlog_system_table.h b/src/paimon/core/table/system/binlog_system_table.h index cfa8f8e62..746ca1511 100644 --- a/src/paimon/core/table/system/binlog_system_table.h +++ b/src/paimon/core/table/system/binlog_system_table.h @@ -26,6 +26,8 @@ namespace paimon { class FileSystem; class TableSchema; +std::shared_ptr CreateBinlogBatchConverter(); + /// System table for `T$binlog`, exposing changelog records with list-wrapped data columns. class BinlogSystemTable : public AuditLogSystemTable { public: diff --git a/src/paimon/core/table/system/metadata_system_tables.cpp b/src/paimon/core/table/system/metadata_system_tables.cpp index c135bfe58..fb17500a7 100644 --- a/src/paimon/core/table/system/metadata_system_tables.cpp +++ b/src/paimon/core/table/system/metadata_system_tables.cpp @@ -41,7 +41,6 @@ #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/field_type_utils.h" #include "paimon/common/utils/internal_row_utils.h" -#include "paimon/common/utils/object_utils.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/rapidjson_util.h" #include "paimon/core/casting/cast_executor_factory.h" @@ -133,7 +132,7 @@ Result> OptionalLocalDateTimePartsToTimestampMillis( return std::optional(timestamp_millis); } -std::optional OptionalDoubleToString(const std::optional& value) { +std::optional OptionalDoubleToString(const std::optional& value) { if (!value) { return std::optional(); } @@ -348,6 +347,45 @@ Result> RowValueStrings(const std::vector& f return values; } +struct StatsStringOverrides { + std::map values; + std::map null_counts; +}; + +Result OmittedPartitionStats(const std::shared_ptr& table_schema, + const DataFileMeta& file, + const BinaryRow& partition, int64_t row_count) { + StatsStringOverrides overrides; + if (!file.write_cols) { + return overrides; + } + + std::vector partition_fields; + partition_fields.reserve(table_schema->PartitionKeys().size()); + for (const auto& partition_key : table_schema->PartitionKeys()) { + PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(partition_key)); + partition_fields.push_back(std::move(field)); + } + PAIMON_ASSIGN_OR_RAISE(std::vector partition_values, + RowValueStrings(partition_fields, partition)); + if (partition_fields.size() != partition_values.size()) { + return Status::Invalid( + fmt::format("partition field count {} does not match partition value count {}", + partition_fields.size(), partition_values.size())); + } + for (size_t i = 0; i < partition_fields.size(); ++i) { + const std::string& field_name = partition_fields[i].Name(); + if (std::find(file.write_cols->begin(), file.write_cols->end(), field_name) != + file.write_cols->end()) { + continue; + } + overrides.values.emplace(field_name, std::move(partition_values[i])); + overrides.null_counts.emplace(field_name, + partition.IsNullAt(i) ? std::to_string(row_count) : "0"); + } + return overrides; +} + Result RowValuesString(const std::vector& fields, const InternalRow& row, std::string_view left, std::string_view right) { PAIMON_ASSIGN_OR_RAISE(std::vector values, RowValueStrings(fields, row)); @@ -366,13 +404,16 @@ Result> OptionalRowValuesString(const std::vector FieldsValueMapString(const std::vector& fields, - const InternalRow& row) { + const InternalRow& row, + const std::map& overrides) { PAIMON_ASSIGN_OR_RAISE(std::vector values, RowValueStrings(fields, row)); std::vector> field_values; size_t length = std::min(fields.size(), values.size()); field_values.reserve(length); for (size_t i = 0; i < length; ++i) { - field_values.emplace_back(fields[i].Name(), std::move(values[i])); + auto iter = overrides.find(fields[i].Name()); + field_values.emplace_back(fields[i].Name(), + iter == overrides.end() ? std::move(values[i]) : iter->second); } std::sort(field_values.begin(), field_values.end(), [](const auto& lhs, const auto& rhs) { return lhs.first < rhs.first; }); @@ -386,13 +427,17 @@ Result FieldsValueMapString(const std::vector& fields, } Result NullValueCountsString(const std::vector& fields, - const InternalArray& null_counts) { + const InternalArray& null_counts, + const std::map& overrides) { std::vector> field_values; int32_t length = std::min(static_cast(fields.size()), null_counts.Size()); field_values.reserve(length); for (int32_t i = 0; i < length; ++i) { + auto iter = overrides.find(fields[i].Name()); std::string value = - null_counts.IsNullAt(i) ? "null" : std::to_string(null_counts.GetLong(i)); + iter == overrides.end() + ? (null_counts.IsNullAt(i) ? "null" : std::to_string(null_counts.GetLong(i))) + : iter->second; field_values.emplace_back(fields[i].Name(), std::move(value)); } std::sort(field_values.begin(), field_values.end(), @@ -422,7 +467,7 @@ Result> ProjectWriteFields(const std::shared_ptr fields; - fields.reserve(file.write_cols->size() + data_schema->PartitionKeys().size()); + fields.reserve(file.write_cols->size()); for (const auto& write_col : file.write_cols.value()) { if (SpecialFields::IsSystemField(write_col)) { continue; @@ -430,15 +475,6 @@ Result> ProjectWriteFields(const std::shared_ptrGetField(write_col)); fields.push_back(std::move(field)); } - - // Partial writes may omit partition columns from write_cols. Keep them in the stats source - // fields so SimpleStatsEvolution can map partition stats consistently. - for (const auto& partition_key : data_schema->PartitionKeys()) { - if (!ObjectUtils::Contains(file.write_cols.value(), partition_key)) { - PAIMON_ASSIGN_OR_RAISE(DataField field, data_schema->GetField(partition_key)); - fields.push_back(std::move(field)); - } - } return fields; } @@ -873,6 +909,9 @@ Result> FilesSystemTable::BuildRows() const { PAIMON_ASSIGN_OR_RAISE( SimpleStatsEvolution::EvolutionStats stats, stats_evolution->Evolution(file->value_stats, file->row_count, file->value_stats_cols)); + PAIMON_ASSIGN_OR_RAISE(StatsStringOverrides stats_overrides, + OmittedPartitionStats(context_.table_schema, *file, + entry.Partition(), file->row_count)); GenericRow row(schema->num_fields()); if (context_.table_schema->PartitionKeys().empty()) { @@ -898,13 +937,16 @@ Result> FilesSystemTable::BuildRows() const { row.SetField(8, OptionalStringValue(min_key)); row.SetField(9, OptionalStringValue(max_key)); PAIMON_ASSIGN_OR_RAISE(std::string null_value_counts, - NullValueCountsString(value_stats_fields, *stats.null_counts)); + NullValueCountsString(value_stats_fields, *stats.null_counts, + stats_overrides.null_counts)); row.SetField(10, StringValue(null_value_counts)); - PAIMON_ASSIGN_OR_RAISE(std::string min_value_stats, - FieldsValueMapString(value_stats_fields, *stats.min_values)); + PAIMON_ASSIGN_OR_RAISE( + std::string min_value_stats, + FieldsValueMapString(value_stats_fields, *stats.min_values, stats_overrides.values)); row.SetField(11, StringValue(min_value_stats)); - PAIMON_ASSIGN_OR_RAISE(std::string max_value_stats, - FieldsValueMapString(value_stats_fields, *stats.max_values)); + PAIMON_ASSIGN_OR_RAISE( + std::string max_value_stats, + FieldsValueMapString(value_stats_fields, *stats.max_values, stats_overrides.values)); row.SetField(12, StringValue(max_value_stats)); row.SetField(13, file->min_sequence_number); row.SetField(14, file->max_sequence_number); diff --git a/src/paimon/core/table/system/system_table_test.cpp b/src/paimon/core/table/system/system_table_test.cpp index c610b8121..fedd7416a 100644 --- a/src/paimon/core/table/system/system_table_test.cpp +++ b/src/paimon/core/table/system/system_table_test.cpp @@ -20,9 +20,11 @@ #include #include #include +#include #include #include "arrow/api.h" +#include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/table/system/audit_log_system_table.h" @@ -31,8 +33,12 @@ #include "paimon/defs.h" #include "paimon/fs/file_system.h" #include "paimon/fs/file_system_factory.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/reader/batch_reader.h" #include "paimon/result.h" #include "paimon/status.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -68,6 +74,92 @@ TEST(SystemTableTest, TestChangelogArrowSchemaReturnsInvalidOptions) { "Invalid Config [table-read.sequence-number.enabled: invalid]"); } +TEST(SystemTableTest, TestBinlogArrowSchemaWithSequenceNumber) { + std::map options = { + {Options::TABLE_READ_SEQUENCE_NUMBER_ENABLED, "true"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + CreateTableSchemaForTest(options)); + + BinlogSystemTable binlog(/*fs=*/nullptr, "/tmp/table", table_schema, options); + ASSERT_OK_AND_ASSIGN(std::shared_ptr schema, binlog.ArrowSchema()); + + ASSERT_EQ(schema->field_names(), + (std::vector{"rowkind", "_SEQUENCE_NUMBER", "pk", "v"})); + ASSERT_EQ(schema->field(0)->type()->id(), arrow::Type::STRING); + ASSERT_FALSE(schema->field(0)->nullable()); + ASSERT_EQ(schema->field(1)->type()->id(), arrow::Type::INT64); + ASSERT_EQ(schema->field(2)->type()->id(), arrow::Type::LIST); + ASSERT_EQ(schema->field(3)->type()->id(), arrow::Type::LIST); +} + +TEST(SystemTableTest, TestStreamingBinlogPacksUpdateAcrossBatches) { + std::shared_ptr input_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("_SEQUENCE_NUMBER", arrow::int64()), + arrow::field("pk", arrow::utf8()), + arrow::field("v", arrow::int32()), + }); + std::shared_ptr input = + arrow::ipc::internal::json::ArrayFromJSON( + input_type, R"([[0, 10, "a", 1], [1, 11, "b", 2], [2, 12, "b", 3], [3, 13, "d", 4]])") + .ValueOrDie(); + std::shared_ptr output_schema = arrow::schema({ + arrow::field("rowkind", arrow::utf8(), /*nullable=*/false), + arrow::field("_SEQUENCE_NUMBER", arrow::int64()), + arrow::field("pk", arrow::list(arrow::utf8())), + arrow::field("v", arrow::list(arrow::int32())), + }); + std::unique_ptr reader = CreateChangelogBatchReader( + std::make_unique(input, input_type, /*read_batch_size=*/2), + output_schema, + /*include_sequence_number=*/true, CreateBinlogBatchConverter(), + /*pack_update_before_after=*/true, GetDefaultPool()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + ReadResultCollector::CollectResult(reader.get())); + std::shared_ptr expected_array = + arrow::ipc::internal::json::ArrayFromJSON(actual->type(), R"([ + ["+I", 10, ["a"], [1]], + ["+U", 11, ["b", "b"], [2, 3]], + ["-D", 13, ["d"], [4]] + ])") + .ValueOrDie(); + auto expected = std::make_shared(expected_array); + ASSERT_TRUE(actual->Equals(*expected)) + << "expected: " << expected->ToString() << "\nactual: " << actual->ToString(); +} + +TEST(SystemTableTest, TestStreamingBinlogEmitsUnmatchedUpdateBefore) { + std::shared_ptr input_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("pk", arrow::utf8()), + arrow::field("v", arrow::int32()), + }); + std::shared_ptr input = + arrow::ipc::internal::json::ArrayFromJSON(input_type, R"([[1, "b", 2]])").ValueOrDie(); + std::shared_ptr output_schema = arrow::schema({ + arrow::field("rowkind", arrow::utf8(), /*nullable=*/false), + arrow::field("pk", arrow::list(arrow::utf8())), + arrow::field("v", arrow::list(arrow::int32())), + }); + std::unique_ptr reader = CreateChangelogBatchReader( + std::make_unique(input, input_type, /*read_batch_size=*/1), + output_schema, + /*include_sequence_number=*/false, CreateBinlogBatchConverter(), + /*pack_update_before_after=*/true, GetDefaultPool()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + ReadResultCollector::CollectResult(reader.get())); + std::shared_ptr expected_array = + arrow::ipc::internal::json::ArrayFromJSON(actual->type(), R"([ + ["-U", ["b"], [2]] + ])") + .ValueOrDie(); + auto expected = std::make_shared(expected_array); + ASSERT_TRUE(actual->Equals(*expected)) + << "expected: " << expected->ToString() << "\nactual: " << actual->ToString(); +} + TEST(SystemTableTest, TestReadOptimizedSystemTableRegistration) { ASSERT_TRUE(SystemTableLoader::IsSupported(ReadOptimizedSystemTable::kName)); diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 78be0ff40..245464a29 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp @@ -36,6 +36,7 @@ #include "gtest/gtest.h" #include "paimon/catalog/catalog.h" #include "paimon/catalog/identifier.h" +#include "paimon/commit_context.h" #include "paimon/common/data/binary_row.h" #include "paimon/common/factories/io_hook.h" #include "paimon/common/reader/complete_row_kind_batch_reader.h" @@ -58,6 +59,8 @@ #include "paimon/data/decimal.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" +#include "paimon/file_store_commit.h" +#include "paimon/file_store_write.h" #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" @@ -78,6 +81,7 @@ #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/test_helper.h" #include "paimon/testing/utils/testharness.h" +#include "paimon/write_context.h" namespace paimon::test { @@ -273,10 +277,12 @@ std::vector StructFieldNames(const std::shared_ptrtype()->fields())->field_names(); } -Result ReadSystemTable( - const std::string& system_table_path, const std::map& options, - bool streaming_mode = false, const std::shared_ptr& predicate = nullptr, - const std::vector& read_field_names = {}) { +Result ReadSystemTable(const std::string& system_table_path, + const std::map& options, + bool streaming_mode = false, + const std::shared_ptr& predicate = nullptr, + const std::vector& read_field_names = {}, + bool read_next_plan = false) { ScanContextBuilder scan_context_builder(system_table_path); scan_context_builder.SetOptions(options).WithStreamingMode(streaming_mode); if (predicate) { @@ -287,6 +293,9 @@ Result ReadSystemTable( PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_scan, TableScan::Create(std::move(scan_context))); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, table_scan->CreatePlan()); + if (read_next_plan) { + PAIMON_ASSIGN_OR_RAISE(plan, table_scan->CreatePlan()); + } ReadContextBuilder read_context_builder(system_table_path); read_context_builder.SetOptions(options); @@ -1207,6 +1216,76 @@ TEST(SystemTableReadInteTest, TestReadFilesSystemTableForPartitionedTable) { ASSERT_EQ(max_value_stats_array->GetString(0), "{dt=20260527, pk=a, v=1}"); } +TEST(SystemTableReadInteTest, TestReadFilesSystemTableForPartitionedPartialWrite) { + arrow::FieldVector fields = { + arrow::field("dt", arrow::utf8()), + arrow::field("id", arrow::int32()), + arrow::field("score", arrow::int32()), + }; + auto schema = arrow::schema(fields); + std::map options = { + {Options::FILE_SYSTEM, "local"}, {Options::FILE_FORMAT, "parquet"}, + {Options::MANIFEST_FORMAT, "avro"}, {Options::BUCKET, "-1"}, + {Options::ROW_TRACKING_ENABLED, "true"}, {Options::DATA_EVOLUTION_ENABLED, "true"}, + }; + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(dir->Str(), schema, + /*partition_keys=*/{"dt"}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/true)); + std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + WriteContextBuilder write_context_builder(table_path, "partial-write"); + write_context_builder.WithWriteSchema({"score"}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, + write_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write, + FileStoreWrite::Create(std::move(write_context))); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_({fields[2]}), R"([[10], [20]])", + /*partition_map=*/{{"dt", "20260724"}}, /*bucket=*/0, {})); + ASSERT_OK(write->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector> commit_messages, + write->PrepareCommit()); + ASSERT_OK(write->Close()); + + CommitContextBuilder commit_context_builder(table_path, "partial-write"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, + commit_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_OK(commit->Commit(commit_messages)); + + ASSERT_OK_AND_ASSIGN(auto files_result, ReadSystemTable(table_path + "$files", options)); + auto files_array = SingleStructChunk(files_result); + ASSERT_EQ(files_array->length(), 1); + auto partition_array = std::dynamic_pointer_cast(files_array->field(0)); + auto null_value_counts_array = + std::dynamic_pointer_cast(files_array->field(10)); + auto min_value_stats_array = + std::dynamic_pointer_cast(files_array->field(11)); + auto max_value_stats_array = + std::dynamic_pointer_cast(files_array->field(12)); + auto write_cols_array = std::dynamic_pointer_cast(files_array->field(19)); + ASSERT_TRUE(partition_array); + ASSERT_TRUE(null_value_counts_array); + ASSERT_TRUE(min_value_stats_array); + ASSERT_TRUE(max_value_stats_array); + ASSERT_TRUE(write_cols_array); + + ASSERT_EQ(partition_array->GetString(0), "{20260724}"); + ASSERT_EQ(null_value_counts_array->GetString(0), "{dt=0, id=2, score=0}"); + ASSERT_EQ(min_value_stats_array->GetString(0), "{dt=20260724, id=null, score=10}"); + ASSERT_EQ(max_value_stats_array->GetString(0), "{dt=20260724, id=null, score=20}"); + auto write_cols_values = + std::dynamic_pointer_cast(write_cols_array->values()); + ASSERT_TRUE(write_cols_values); + ASSERT_EQ(write_cols_values->length(), 1); + ASSERT_EQ(write_cols_values->GetString(0), "score"); +} + TEST(SystemTableReadInteTest, TestReadFilesSystemTableForDatePartition) { arrow::FieldVector fields = { arrow::field("dt", arrow::date32()), @@ -1610,6 +1689,67 @@ TEST(SystemTableReadInteTest, TestReadAuditLogAndBinlogSystemTableWithChangelogR ])"); } +TEST(SystemTableReadInteTest, TestStreamingBinlogPacksUpdateBeforeAndAfter) { + arrow::FieldVector fields = { + arrow::field("pk", arrow::utf8()), + arrow::field("v", arrow::int32()), + }; + auto schema = arrow::schema(fields); + std::map options = { + {Options::FILE_SYSTEM, "local"}, {Options::FILE_FORMAT, "parquet"}, + {Options::MANIFEST_FORMAT, "avro"}, {Options::BUCKET, "1"}, + {Options::WRITE_BUFFER_SIZE, "1"}, {Options::WRITE_BUFFER_SPILLABLE, "false"}, + }; + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(dir->Str(), schema, + /*partition_keys=*/{}, + /*primary_keys=*/{"pk"}, options, + /*is_streaming_mode=*/true)); + + std::vector row_kinds_1 = { + RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_BEFORE, + }; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch_1, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["a", 1], ["b", 2]])", + /*partition_map=*/{}, /*bucket=*/0, row_kinds_1)); + std::vector row_kinds_2 = { + RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, + }; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch_2, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["b", 3], ["c", 4]])", + /*partition_map=*/{}, /*bucket=*/0, row_kinds_2)); + std::vector> batches; + batches.push_back(std::move(batch_1)); + batches.push_back(std::move(batch_2)); + ASSERT_OK(helper->WriteAndCommit(std::move(batches), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + std::map streaming_options = options; + streaming_options[Options::SCAN_MODE] = "from-snapshot"; + streaming_options[Options::SCAN_SNAPSHOT_ID] = "1"; + ASSERT_OK_AND_ASSIGN( + auto result, + ReadSystemTable(PathUtil::JoinPath(dir->Str(), "foo.db/bar$binlog"), streaming_options, + /*streaming_mode=*/true, /*predicate=*/nullptr, + /*read_field_names=*/{}, /*read_next_plan=*/true)); + ASSERT_TRUE(result.array); + ASSERT_EQ(result.array->num_chunks(), 2); + auto array = std::dynamic_pointer_cast( + arrow::Concatenate(result.array->chunks()).ValueOrDie()); + ASSERT_TRUE(array); + ASSERT_EQ(StructFieldNames(array), (std::vector{"rowkind", "pk", "v"})); + AssertStructArrayEqualsJson(array, R"([ + ["+I", ["a"], [1]], + ["+U", ["b", "b"], [2, 3]], + ["-D", ["c"], [4]] + ])"); +} + TEST(SystemTableReadInteTest, TestReadBinlogSystemTableWithNullValue) { arrow::FieldVector fields = { arrow::field("pk", arrow::utf8(), /*nullable=*/false),