From 1426314b902e7394f69e3ad016be67c03567e073 Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Sat, 29 Aug 2026 18:03:09 +0800 Subject: [PATCH 1/8] feat(read): decouple read batch lifetime from readers --- include/paimon/format/reader_builder.h | 4 + include/paimon/reader/batch_reader.h | 12 ++ include/paimon/table/source/table_read.h | 14 +- ...hared_shredding_read_plan_factory_test.cpp | 41 ++-- .../data/shredding/shredding_file_reader.cpp | 5 +- .../data/shredding/shredding_file_reader.h | 2 +- .../apply_bitmap_index_batch_reader_test.cpp | 16 +- .../complete_index_score_batch_reader.cpp | 10 +- .../complete_index_score_batch_reader.h | 4 +- ...complete_index_score_batch_reader_test.cpp | 16 +- .../reader/blob_fallback_batch_reader.cpp | 20 +- .../reader/blob_fallback_batch_reader.h | 4 +- .../blob_fallback_batch_reader_test.cpp | 31 +-- .../blob_view_resolving_batch_reader.cpp | 6 +- .../reader/blob_view_resolving_batch_reader.h | 5 +- .../blob_view_resolving_batch_reader_test.cpp | 21 ++- .../reader/complete_row_kind_batch_reader.cpp | 7 +- .../reader/complete_row_kind_batch_reader.h | 6 +- .../complete_row_kind_batch_reader_test.cpp | 18 +- .../common/reader/concat_batch_reader.cpp | 9 +- .../common/reader/concat_batch_reader.h | 4 +- .../reader/concat_batch_reader_test.cpp | 8 +- .../reader/data_evolution_file_reader.cpp | 5 +- .../reader/data_evolution_file_reader.h | 2 +- .../data_evolution_file_reader_test.cpp | 39 ++-- .../late_materializing_file_batch_reader.cpp | 15 +- .../late_materializing_file_batch_reader.h | 3 +- ...e_materializing_file_batch_reader_test.cpp | 98 ++++++---- .../late_materializing_reader_builder.h | 9 +- .../common/reader/predicate_batch_reader.cpp | 13 +- .../common/reader/predicate_batch_reader.h | 6 +- .../reader/predicate_batch_reader_test.cpp | 12 +- .../prefetch_file_batch_reader_impl.cpp | 10 +- .../reader/prefetch_file_batch_reader_impl.h | 7 +- .../prefetch_file_batch_reader_impl_test.cpp | 78 ++++---- src/paimon/common/reader/reader_utils.cpp | 7 +- src/paimon/common/reader/reader_utils.h | 3 +- .../common/reader/reader_utils_test.cpp | 60 +++++- src/paimon/common/utils/arrow/mem_utils.cpp | 46 +++-- src/paimon/common/utils/arrow/mem_utils.h | 13 +- .../common/utils/arrow/mem_utils_test.cpp | 143 ++++++++++++++ .../core/append/append_only_writer_test.cpp | 6 +- ...pply_deletion_vector_batch_reader_test.cpp | 5 +- .../io/async_key_value_projection_reader.h | 7 +- .../complete_row_tracking_fields_reader.cpp | 5 +- .../io/complete_row_tracking_fields_reader.h | 2 +- ...mplete_row_tracking_fields_reader_test.cpp | 18 +- src/paimon/core/io/field_mapping_reader.cpp | 9 +- src/paimon/core/io/field_mapping_reader.h | 4 +- .../core/io/field_mapping_reader_test.cpp | 28 +-- .../generic_row_to_arrow_array_converter.cpp | 6 +- .../io/generic_row_to_arrow_array_converter.h | 7 +- .../io/key_value_meta_projection_consumer.cpp | 4 +- .../io/key_value_meta_projection_consumer.h | 4 +- .../core/io/key_value_projection_consumer.cpp | 6 +- .../core/io/key_value_projection_consumer.h | 7 +- .../core/io/key_value_projection_reader.cpp | 7 +- .../core/io/key_value_projection_reader.h | 2 +- .../io/key_value_projection_reader_test.cpp | 18 +- .../core/io/meta_to_arrow_array_converter.cpp | 4 +- .../core/io/meta_to_arrow_array_converter.h | 4 +- .../core/io/row_to_arrow_array_converter.h | 9 +- .../core/io/vector_file_batch_reader.cpp | 5 +- src/paimon/core/io/vector_file_batch_reader.h | 2 +- .../core/io/vector_file_batch_reader_test.cpp | 13 +- .../lookup_merge_tree_compact_rewriter.cpp | 4 +- ...ookup_merge_tree_compact_rewriter_test.cpp | 2 +- .../compact/merge_tree_compact_rewriter.cpp | 4 +- .../merge_tree_compact_rewriter_test.cpp | 2 +- src/paimon/core/mergetree/lookup_levels.cpp | 4 +- .../core/mergetree/merge_tree_writer_test.cpp | 2 +- .../core/operation/abstract_split_read.cpp | 28 +-- .../core/operation/abstract_split_read.h | 2 + .../append_only_file_store_write.cpp | 7 +- .../append_only_file_store_write_test.cpp | 2 +- .../operation/data_evolution_split_read.cpp | 25 +-- .../core/operation/internal_read_context.cpp | 18 +- .../core/operation/internal_read_context.h | 11 +- .../operation/internal_read_context_test.cpp | 55 +++--- .../core/operation/merge_file_split_read.cpp | 18 +- .../operation/merge_file_split_read_test.cpp | 36 ++-- .../core/operation/raw_file_split_read.cpp | 5 +- .../operation/raw_file_split_read_test.cpp | 39 ++-- .../postpone/postpone_bucket_writer_test.cpp | 2 +- .../core/realtime/arrow_realtime_store.cpp | 5 +- .../realtime/arrow_realtime_store_factory.cpp | 3 +- .../realtime/realtime_append_only_writer.cpp | 5 +- .../realtime/realtime_append_only_writer.h | 3 +- .../table/source/append_only_table_read.cpp | 16 +- .../table/source/fallback_data_split_test.cpp | 18 +- .../core/table/source/fallback_table_read.cpp | 12 ++ .../core/table/source/fallback_table_read.h | 16 +- .../table/source/key_value_table_read.cpp | 19 +- .../core/table/source/key_value_table_read.h | 1 - .../table/source/pk_count_reader_test.cpp | 11 +- src/paimon/core/table/source/table_read.cpp | 29 +-- .../table/system/audit_log_system_table.cpp | 49 ++--- .../table/system/in_memory_system_table.cpp | 13 +- .../core/table/system/system_table_test.cpp | 41 +++- .../core/utils/manifest_meta_reader.cpp | 5 +- src/paimon/core/utils/manifest_meta_reader.h | 4 +- src/paimon/core/utils/objects_file.h | 5 +- .../format/avro/avro_file_batch_reader.cpp | 10 +- .../format/avro/avro_file_batch_reader.h | 7 +- .../avro/avro_file_batch_reader_test.cpp | 25 ++- .../format/avro/avro_file_format_test.cpp | 2 +- .../format/avro/avro_format_writer_test.cpp | 6 +- src/paimon/format/avro/avro_reader_builder.h | 9 +- .../format/blob/blob_file_batch_reader.cpp | 20 +- .../format/blob/blob_file_batch_reader.h | 6 +- .../blob/blob_file_batch_reader_test.cpp | 75 ++++---- .../format/blob/blob_format_writer_test.cpp | 62 +++--- src/paimon/format/blob/blob_reader_builder.h | 10 +- .../format/blob/blob_stats_extractor.cpp | 4 +- .../mosaic/mosaic_file_batch_reader.cpp | 4 +- .../format/mosaic/mosaic_file_batch_reader.h | 3 +- .../format/mosaic/mosaic_file_format_test.cpp | 2 +- .../format/mosaic/mosaic_reader_builder.h | 9 +- .../format/orc/complex_predicate_test.cpp | 6 +- src/paimon/format/orc/orc_adapter_test.cpp | 29 +-- .../format/orc/orc_file_batch_reader.cpp | 22 +-- src/paimon/format/orc/orc_file_batch_reader.h | 22 ++- .../format/orc/orc_file_batch_reader_test.cpp | 129 +++++++++---- src/paimon/format/orc/orc_reader_builder.h | 10 +- .../format/orc/predicate_pushdown_test.cpp | 8 +- .../page_filtered_row_group_reader_test.cpp | 8 +- .../parquet/parquet_file_batch_reader.cpp | 2 + .../parquet/parquet_file_batch_reader.h | 4 +- .../parquet_file_batch_reader_test.cpp | 42 ++--- .../format/parquet/parquet_reader_builder.h | 14 +- .../format/parquet/parquet_vector_io_test.cpp | 25 ++- .../parquet/predicate_pushdown_test.cpp | 4 +- .../format/parquet/variant_parquet_test.cpp | 12 +- .../testing/mock/mock_file_batch_reader.h | 5 +- .../testing/utils/read_result_collector.h | 93 ++++++--- src/paimon/testing/utils/test_helper.h | 13 +- test/inte/append_compaction_inte_test.cpp | 3 +- test/inte/blob_table_inte_test.cpp | 8 +- test/inte/data_evolution_table_test.cpp | 20 +- test/inte/global_index_test.cpp | 2 +- test/inte/nested_column_pruning_inte_test.cpp | 9 +- test/inte/pk_compaction_inte_test.cpp | 15 +- .../primary_key_sorted_index_inte_test.cpp | 2 +- test/inte/read_inte_test.cpp | 177 ++++++++++-------- test/inte/read_inte_with_index_test.cpp | 21 ++- test/inte/realtime_write_inte_test.cpp | 9 +- test/inte/scan_and_read_inte_test.cpp | 134 ++++++++----- test/inte/write_and_read_inte_test.cpp | 26 ++- test/inte/write_inte_test.cpp | 18 +- 149 files changed, 1670 insertions(+), 1011 deletions(-) diff --git a/include/paimon/format/reader_builder.h b/include/paimon/format/reader_builder.h index 5a28077a8..c9379dbb5 100644 --- a/include/paimon/format/reader_builder.h +++ b/include/paimon/format/reader_builder.h @@ -53,6 +53,10 @@ class PAIMON_EXPORT ReaderBuilder { } /// Build a file batch reader based on the created `InputStream`. + /// + /// Every non-EOF ArrowArray returned by the reader must retain all allocator and plugin + /// resources needed by its release callback. The array must remain releasable after the + /// reader has been destroyed. virtual Result> Build( const std::shared_ptr& path) const = 0; }; diff --git a/include/paimon/reader/batch_reader.h b/include/paimon/reader/batch_reader.h index c0033c206..87752a64e 100644 --- a/include/paimon/reader/batch_reader.h +++ b/include/paimon/reader/batch_reader.h @@ -45,6 +45,12 @@ class PAIMON_EXPORT BatchReader { /// retried, as it will repeatedly return the same error code. /// \note IMPORTANT: A non-EOF ArrowArray and all its nested child arrays must have offset 0 to /// avoid potential issues during conversion through the Arrow C Data Interface. + /// \note A returned ArrowArray must retain every allocator and plugin resource needed by its + /// release callback, so it remains releasable after this reader is destroyed. + /// \note Consumers must treat the returned ArrowArray and ArrowSchema as one complete Arrow C + /// Data Interface ownership unit. Moving or retaining an individual child ArrowArray without + /// its root array is unsupported because resource lifetimes are retained by the root array's + /// release chain. /// /// @return A result containing a `::ReadBatch`, which consists of a unique pointer to /// `ArrowArray` and a unique pointer to `ArrowSchema`. Returned array contains a `_VALUE_KIND` @@ -59,6 +65,12 @@ class PAIMON_EXPORT BatchReader { /// retried, as it will repeatedly return the same error code. /// \note IMPORTANT: A non-EOF ArrowArray and all its nested child arrays must have offset 0 to /// avoid potential issues during conversion through the Arrow C Data Interface. + /// \note A returned ArrowArray must retain every allocator and plugin resource needed by its + /// release callback, so it remains releasable after this reader is destroyed. + /// \note Consumers must treat the returned ArrowArray and ArrowSchema as one complete Arrow C + /// Data Interface ownership unit. Moving or retaining an individual child ArrowArray without + /// its root array is unsupported because resource lifetimes are retained by the root array's + /// release chain. /// /// @return A result containing a `::ReadBatch` and a valid bitmap. `::ReadBatch` consists of a /// unique pointer to `ArrowArray` and a unique pointer to `ArrowSchema`. Returned array diff --git a/include/paimon/table/source/table_read.h b/include/paimon/table/source/table_read.h index c9cfe7c85..c77b323da 100644 --- a/include/paimon/table/source/table_read.h +++ b/include/paimon/table/source/table_read.h @@ -23,7 +23,6 @@ #include #include "paimon/executor.h" -#include "paimon/memory/memory_pool.h" #include "paimon/read_context.h" #include "paimon/reader/batch_reader.h" #include "paimon/reader/count_reader.h" @@ -32,7 +31,6 @@ #include "paimon/visibility.h" namespace paimon { -class MemoryPool; class ReadContext; /// Given a `Split` or a list of `Split`, generate a reader for batch reading. @@ -61,7 +59,7 @@ class PAIMON_EXPORT TableRead { /// @note `BatchReader`s created by the same `TableRead` are not thread-safe for /// concurrent reading. virtual Result> CreateReader( - const std::vector>& splits); + const std::vector>& splits) = 0; /// Creates a `BatchReader` instance for a single split. /// @@ -76,15 +74,5 @@ class PAIMON_EXPORT TableRead { /// Implementations may override this to provide a more efficient count path. virtual Result> CreateCountReader( const std::vector>& splits); - - protected: - explicit TableRead(const std::shared_ptr& memory_pool); - - std::shared_ptr GetMemoryPool() const { - return pool_; - } - - private: - std::shared_ptr pool_; }; } // namespace paimon diff --git a/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp index e02f0ffb0..47e9649bd 100644 --- a/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp +++ b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp @@ -35,6 +35,7 @@ #include "paimon/common/data/shredding/map_shredding_defs.h" #include "paimon/common/data/shredding/shredding_file_reader.h" #include "paimon/common/fs/external_path_provider.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/append/append_only_writer.h" #include "paimon/core/compact/noop_compact_manager.h" @@ -56,6 +57,7 @@ class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } static MapSharedShreddingFieldMeta TagsMeta() { @@ -137,7 +139,7 @@ class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test { field_read_plans.emplace(field->name(), std::move(field_read_plan)); } return std::make_unique(std::move(reader), std::move(field_read_plans), - pool_); + arrow_pool_); } Result> CreateReader( @@ -224,6 +226,7 @@ class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test { private: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::shared_ptr logical_schema_ = arrow::schema({ arrow::field("id", arrow::int32()), arrow::field("tags", arrow::map(arrow::utf8(), arrow::int64())), @@ -246,7 +249,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestAllExistSelectedKeysWithoutOve auto read_schema = ExportSchema(ReadSchema("b")); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -268,7 +271,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestAllExistSelectedKeysWithOverfl auto read_schema = ExportSchema(ReadSchema("a,c")); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -301,14 +304,14 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjection) selected_field, TagsMeta())); std::map> contexts; contexts.emplace("tags", std::move(field_read_plan)); - auto reader = - std::make_unique(std::move(mock_reader), std::move(contexts), pool_); + auto reader = std::make_unique(std::move(mock_reader), std::move(contexts), + arrow_pool_); auto read_schema = ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field})); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); auto expected_type = arrow::struct_({arrow::field("id", arrow::int32()), selected_field}); std::shared_ptr expected; @@ -441,14 +444,14 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjectionFr file_schema->field(1), selected_field)); std::map> contexts; contexts.emplace("tags", std::move(field_read_plan)); - auto reader = - std::make_unique(std::move(mock_reader), std::move(contexts), pool_); + auto reader = std::make_unique(std::move(mock_reader), std::move(contexts), + arrow_pool_); auto read_schema = ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field})); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); auto expected_type = arrow::struct_({arrow::field("id", arrow::int32()), selected_field}); std::shared_ptr expected; @@ -494,7 +497,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestPartialExistSelectedKeys) { ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -514,7 +517,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestMissingSelectedKeysReadsWholeM auto read_schema = ExportSchema(ReadSchema(std::nullopt)); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -551,7 +554,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSpecialSelectedKeys) { auto read_schema = ExportSchema(ReadSchema(selected_keys)); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -590,7 +593,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestUnknownSelectedKeyReturnsEmpty ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -618,7 +621,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestInvalidNullFieldMappingField) auto read_schema = ExportSchema(ReadSchema("a")); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_NOK_WITH_MSG(ReadResultCollector::CollectResult(reader.get()), + ASSERT_NOK_WITH_MSG(ReadResultCollector::CollectResult(std::move(reader)), "__field_mapping cannot be null"); } @@ -635,7 +638,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestInvalidNullFieldMappingFieldEl auto read_schema = ExportSchema(ReadSchema("b")); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_NOK_WITH_MSG(ReadResultCollector::CollectResult(reader.get()), + ASSERT_NOK_WITH_MSG(ReadResultCollector::CollectResult(std::move(reader)), "__field_mapping element cannot be null"); } @@ -680,7 +683,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestListValue) { auto read_schema = ExportSchema(arrow::schema(std::move(read_fields))); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -741,7 +744,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestOrcDictionaryEncodedStringValu auto read_schema = ExportSchema(arrow::schema(std::move(read_fields))); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( arrow::struct_(logical_schema->fields()), {R"([ @@ -801,7 +804,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestOrcDictionaryEncodedStringList auto read_schema = ExportSchema(arrow::schema(std::move(read_fields))); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( arrow::struct_(logical_schema->fields()), {R"([ @@ -852,7 +855,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestReadsRealFormatFile) { auto read_schema = ExportSchema(ReadSchema("a,c")); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( diff --git a/src/paimon/common/data/shredding/shredding_file_reader.cpp b/src/paimon/common/data/shredding/shredding_file_reader.cpp index 1c2e34f57..eb22ab33a 100644 --- a/src/paimon/common/data/shredding/shredding_file_reader.cpp +++ b/src/paimon/common/data/shredding/shredding_file_reader.cpp @@ -33,8 +33,8 @@ namespace paimon { ShreddingFileReader::ShreddingFileReader( std::unique_ptr&& reader, std::map>&& plans, - const std::shared_ptr& pool) - : arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)), plans_(std::move(plans)) {} + const std::shared_ptr& arrow_pool) + : arrow_pool_(arrow_pool), reader_(std::move(reader)), plans_(std::move(plans)) {} Result> ShreddingFileReader::GetFileSchema() const { return reader_->GetFileSchema(); @@ -109,6 +109,7 @@ Result ShreddingFileReader::NextBatchWithBitma auto new_c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*new_struct_array, new_c_array.get(), new_c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(new_c_array.get(), arrow_pool_)); batch = std::make_pair(std::move(new_c_array), std::move(new_c_schema)); return batch_with_bitmap; } diff --git a/src/paimon/common/data/shredding/shredding_file_reader.h b/src/paimon/common/data/shredding/shredding_file_reader.h index f6ab35759..d3584c7db 100644 --- a/src/paimon/common/data/shredding/shredding_file_reader.h +++ b/src/paimon/common/data/shredding/shredding_file_reader.h @@ -38,7 +38,7 @@ class ShreddingFileReader : public FileBatchReader { public: ShreddingFileReader(std::unique_ptr&& reader, std::map>&& plans, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result> GetFileSchema() const override; diff --git a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp index ea45c9791..cd713ab83 100644 --- a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp +++ b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp @@ -29,6 +29,7 @@ #include "fmt/ranges.h" #include "gtest/gtest.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/read_ahead_cache.h" #include "paimon/executor.h" @@ -55,15 +56,17 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, target_type_ = arrow::struct_({arrow::field("f1", int_type_)}); pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); fs_ = std::make_shared(); ASSERT_OK_AND_ASSIGN(executor_, CreateDefaultExecutor(/*thread_count=*/2)); } void TearDown() override {} - void CheckResult(BatchReader* apply_bitmap_batch_reader, + void CheckResult(std::unique_ptr apply_bitmap_batch_reader, const std::shared_ptr& expected_chunk_array) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr result_chunk_array, - ReadResultCollector::CollectResult(apply_bitmap_batch_reader)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr result_chunk_array, + ReadResultCollector::CollectResult(std::move(apply_bitmap_batch_reader))); if (expected_chunk_array) { ASSERT_TRUE(result_chunk_array); ASSERT_EQ(expected_chunk_array->length(), result_chunk_array->length()); @@ -98,7 +101,7 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_)); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); @@ -106,14 +109,14 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, auto apply_bitmap_batch_reader = std::make_unique( std::move(file_batch_reader), std::move(bitmap_index)); if (expected_str.empty()) { - CheckResult(apply_bitmap_batch_reader.get(), nullptr); + CheckResult(std::move(apply_bitmap_batch_reader), nullptr); } else { auto expected = arrow::ipc::internal::json::ArrayFromJSON(int_type_, expected_str).ValueOrDie(); std::shared_ptr expect_array = arrow::StructArray::Make({expected}, target_type_->fields()).ValueOrDie(); auto expected_chunk_array = std::make_shared(expect_array); - CheckResult(apply_bitmap_batch_reader.get(), expected_chunk_array); + CheckResult(std::move(apply_bitmap_batch_reader), expected_chunk_array); } } } @@ -122,6 +125,7 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, std::shared_ptr int_type_; std::shared_ptr target_type_; std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::shared_ptr fs_; std::shared_ptr executor_; }; diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader.cpp b/src/paimon/common/global_index/complete_index_score_batch_reader.cpp index c46263221..65e0c72ee 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader.cpp +++ b/src/paimon/common/global_index/complete_index_score_batch_reader.cpp @@ -38,13 +38,16 @@ namespace paimon { CompleteIndexScoreBatchReader::CompleteIndexScoreBatchReader( std::unique_ptr&& reader, const std::vector& scores, - const std::shared_ptr& pool) - : arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)), scores_(scores) {} + const std::shared_ptr& arrow_pool) + : arrow_pool_(arrow_pool), reader_(std::move(reader)), scores_(scores) {} Result CompleteIndexScoreBatchReader::NextBatch() { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, NextBatchWithBitmap()); - return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool_.get()); + PAIMON_ASSIGN_OR_RAISE( + BatchReader::ReadBatch batch, + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool_)); + return batch; } void CompleteIndexScoreBatchReader::UpdateScoreFieldIndex(const arrow::StructType* struct_type) { @@ -106,6 +109,7 @@ Result CompleteIndexScoreBatchReader::NextBatc arrow::StructArray::Make(array_vec, field_names_with_score_)); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*array_with_score, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return batch_with_bitmap; } } // namespace paimon diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader.h b/src/paimon/common/global_index/complete_index_score_batch_reader.h index ac97b7e9a..427d854bb 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader.h +++ b/src/paimon/common/global_index/complete_index_score_batch_reader.h @@ -42,7 +42,7 @@ class CompleteIndexScoreBatchReader : public BatchReader { public: CompleteIndexScoreBatchReader(std::unique_ptr&& reader, const std::vector& scores, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result NextBatch() override; @@ -63,7 +63,7 @@ class CompleteIndexScoreBatchReader : public BatchReader { size_t score_cursor_ = 0; int32_t index_score_field_idx_ = -1; std::vector field_names_with_score_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::unique_ptr reader_; std::vector scores_; }; diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp b/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp index e14b28f6a..1d93fe2d6 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp +++ b/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp @@ -27,6 +27,7 @@ #include "gtest/gtest.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/file_format.h" #include "paimon/format/file_format_factory.h" #include "paimon/memory/memory_pool.h" @@ -40,6 +41,7 @@ class CompleteIndexScoreBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } void TearDown() override { pool_.reset(); @@ -51,7 +53,7 @@ class CompleteIndexScoreBatchReaderTest : public ::testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), selected_bitmap, batch_size); return std::make_unique(std::move(file_batch_reader), scores, - pool_); + arrow_pool_); } std::unique_ptr PrepareCompleteIndexScoreBatchReader( @@ -60,11 +62,12 @@ class CompleteIndexScoreBatchReaderTest : public ::testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), batch_size); return std::make_unique(std::move(file_batch_reader), scores, - pool_); + arrow_pool_); } private: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; TEST_F(CompleteIndexScoreBatchReaderTest, TestSimple) { @@ -86,7 +89,7 @@ TEST_F(CompleteIndexScoreBatchReaderTest, TestSimple) { auto reader = PrepareCompleteIndexScoreBatchReader(src_array, scores, /*batch_size=*/1); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected_array; auto array_status = @@ -98,7 +101,6 @@ TEST_F(CompleteIndexScoreBatchReaderTest, TestSimple) { &expected_array); ASSERT_TRUE(array_status.ok()); ASSERT_TRUE(expected_array->ApproxEquals(*result_array)); - reader->Close(); } TEST_F(CompleteIndexScoreBatchReaderTest, TestWithBitmap) { @@ -122,7 +124,7 @@ TEST_F(CompleteIndexScoreBatchReaderTest, TestWithBitmap) { auto reader = PrepareCompleteIndexScoreBatchReader(src_array, selected_bitmap, scores, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected_array; auto array_status = @@ -133,7 +135,6 @@ TEST_F(CompleteIndexScoreBatchReaderTest, TestWithBitmap) { &expected_array); ASSERT_TRUE(array_status.ok()); ASSERT_TRUE(expected_array->ApproxEquals(*result_array)); - reader->Close(); } TEST_F(CompleteIndexScoreBatchReaderTest, TestReadWithNullScores) { @@ -155,11 +156,10 @@ TEST_F(CompleteIndexScoreBatchReaderTest, TestReadWithNullScores) { auto reader = PrepareCompleteIndexScoreBatchReader(src_array, /*scores=*/{}, /*batch_size=*/1); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto expected_array = std::make_shared(src_array); ASSERT_TRUE(expected_array->Equals(*result_array)); - reader->Close(); } } // namespace paimon::test diff --git a/src/paimon/common/reader/blob_fallback_batch_reader.cpp b/src/paimon/common/reader/blob_fallback_batch_reader.cpp index 21db0ade6..bdd312238 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader.cpp +++ b/src/paimon/common/reader/blob_fallback_batch_reader.cpp @@ -41,7 +41,7 @@ namespace paimon { Result> BlobFallbackBatchReader::Create( std::vector>&& sequence_groups, const std::shared_ptr& read_schema, int32_t read_batch_size, - const std::shared_ptr& pool) { + const std::shared_ptr& arrow_pool) { if (sequence_groups.size() < 2) { return Status::Invalid( "Blob fallback needs at least two sequence groups; a single group should be read " @@ -85,23 +85,22 @@ Result> BlobFallbackBatchReader::Create cursor.segments = std::move(segments); groups.push_back(std::move(cursor)); } - return std::unique_ptr( - new BlobFallbackBatchReader(std::move(groups), read_schema, blob_field_idx, - row_id_field_idx, seq_num_field_idx, read_batch_size, pool)); + return std::unique_ptr(new BlobFallbackBatchReader( + std::move(groups), read_schema, blob_field_idx, row_id_field_idx, seq_num_field_idx, + read_batch_size, arrow_pool)); } -BlobFallbackBatchReader::BlobFallbackBatchReader(std::vector&& groups, - const std::shared_ptr& read_schema, - int32_t blob_field_idx, int32_t row_id_field_idx, - int32_t seq_num_field_idx, int32_t read_batch_size, - const std::shared_ptr& pool) +BlobFallbackBatchReader::BlobFallbackBatchReader( + std::vector&& groups, const std::shared_ptr& read_schema, + int32_t blob_field_idx, int32_t row_id_field_idx, int32_t seq_num_field_idx, + int32_t read_batch_size, const std::shared_ptr& arrow_pool) : groups_(std::move(groups)), read_schema_(read_schema), blob_field_idx_(blob_field_idx), row_id_field_idx_(row_id_field_idx), seq_num_field_idx_(seq_num_field_idx), read_batch_size_(read_batch_size), - arrow_pool_(GetArrowPool(pool)) {} + arrow_pool_(arrow_pool) {} Result BlobFallbackBatchReader::FillWindow(size_t group_idx, int64_t want, std::vector* chunks) { @@ -365,6 +364,7 @@ Result BlobFallbackBatchReader::NextBatch() { std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*normalized_array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return std::make_pair(std::move(c_array), std::move(c_schema)); } diff --git a/src/paimon/common/reader/blob_fallback_batch_reader.h b/src/paimon/common/reader/blob_fallback_batch_reader.h index bd7902076..1e30b8957 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader.h +++ b/src/paimon/common/reader/blob_fallback_batch_reader.h @@ -81,7 +81,7 @@ class BlobFallbackBatchReader : public BatchReader { static Result> Create( std::vector>&& sequence_groups, const std::shared_ptr& read_schema, int32_t read_batch_size, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result NextBatch() override; @@ -127,7 +127,7 @@ class BlobFallbackBatchReader : public BatchReader { const std::shared_ptr& read_schema, int32_t blob_field_idx, int32_t row_id_field_idx, int32_t seq_num_field_idx, int32_t read_batch_size, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); /// Collects up to `want` rows from the group into chunks. Only the first group may come up /// short (which defines the window size); any later group ending early is a misalignment. diff --git a/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp b/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp index a1ddf9621..1750d8ef4 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp +++ b/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp @@ -29,6 +29,7 @@ #include "paimon/common/data/blob_defs.h" #include "paimon/common/data/blob_utils.h" #include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/mock/mock_file_batch_reader.h" @@ -45,6 +46,7 @@ class BlobFallbackBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); struct_type_ = arrow::struct_({BlobUtils::ToArrowField("blob_col", true)}); read_schema_ = arrow::schema(struct_type_->fields()); } @@ -119,10 +121,9 @@ class BlobFallbackBatchReaderTest : public ::testing::Test { } ASSERT_OK_AND_ASSIGN( auto reader, BlobFallbackBatchReader::Create(std::move(groups), read_schema_, - batch_size, pool_)); - ASSERT_OK_AND_ASSIGN( - auto result, paimon::test::ReadResultCollector::CollectResult(reader.get())); - reader->Close(); + batch_size, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto result, paimon::test::ReadResultCollector::CollectResult( + std::move(reader))); auto expected_chunk_array = std::make_shared(expected_array); ASSERT_TRUE(result->Equals(expected_chunk_array)) << "batch_size=" << batch_size << " file_batch_size=" << file_batch_size @@ -134,6 +135,7 @@ class BlobFallbackBatchReaderTest : public ::testing::Test { protected: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::shared_ptr struct_type_; std::shared_ptr read_schema_; }; @@ -292,11 +294,10 @@ TEST_F(BlobFallbackBatchReaderTest, TestRowTrackingFieldsPreserved) { ASSERT_OK_AND_ASSIGN(auto reader, BlobFallbackBatchReader::Create(std::move(groups), schema, - batch_size, pool_)); + batch_size, arrow_pool_)); ASSERT_OK_AND_ASSIGN( auto result, - paimon::test::ReadResultCollector::CollectResult(reader.get())); - reader->Close(); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); // row 0 falls back to seq 10, row 1 is all-placeholder (null blob, row id // kept, seq -1), row 2 takes seq 20 @@ -319,8 +320,8 @@ TEST_F(BlobFallbackBatchReaderTest, TestMisalignedGroupsFail) { std::vector> groups; groups.push_back(MakeGroup({SegmentSpec::File({"PH", "u1", "PH"})}, 1024)); groups.push_back(MakeGroup({SegmentSpec::File({"b0", "b1"})}, 1024)); - ASSERT_OK_AND_ASSIGN( - auto reader, BlobFallbackBatchReader::Create(std::move(groups), read_schema_, 1024, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, BlobFallbackBatchReader::Create( + std::move(groups), read_schema_, 1024, arrow_pool_)); ASSERT_NOK_WITH_MSG(reader->NextBatch(), "same number of rows"); } @@ -329,7 +330,7 @@ TEST_F(BlobFallbackBatchReaderTest, TestCreateValidation) { std::vector> single_group; single_group.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); ASSERT_NOK_WITH_MSG( - BlobFallbackBatchReader::Create(std::move(single_group), read_schema_, 1024, pool_), + BlobFallbackBatchReader::Create(std::move(single_group), read_schema_, 1024, arrow_pool_), "at least two sequence groups"); // the read schema must contain a blob field @@ -339,23 +340,23 @@ TEST_F(BlobFallbackBatchReaderTest, TestCreateValidation) { auto plain_schema = arrow::schema({arrow::field("not_blob", arrow::large_binary(), /*nullable=*/true)}); ASSERT_NOK_WITH_MSG( - BlobFallbackBatchReader::Create(std::move(groups), plain_schema, 1024, pool_), + BlobFallbackBatchReader::Create(std::move(groups), plain_schema, 1024, arrow_pool_), "should contain a blob field"); // groups must not be empty std::vector> with_empty_group; with_empty_group.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); with_empty_group.emplace_back(); - ASSERT_NOK_WITH_MSG( - BlobFallbackBatchReader::Create(std::move(with_empty_group), read_schema_, 1024, pool_), - "should not be empty"); + ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(with_empty_group), read_schema_, + 1024, arrow_pool_), + "should not be empty"); // a gap segment must cover at least one selected row id std::vector> with_empty_gap; with_empty_gap.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); with_empty_gap.push_back(MakeGroup({SegmentSpec::GapRanges({})}, 1024)); ASSERT_NOK_WITH_MSG( - BlobFallbackBatchReader::Create(std::move(with_empty_gap), read_schema_, 1024, pool_), + BlobFallbackBatchReader::Create(std::move(with_empty_gap), read_schema_, 1024, arrow_pool_), "at least one selected row id"); } diff --git a/src/paimon/common/reader/blob_view_resolving_batch_reader.cpp b/src/paimon/common/reader/blob_view_resolving_batch_reader.cpp index 29adedcd5..822e202d1 100644 --- a/src/paimon/common/reader/blob_view_resolving_batch_reader.cpp +++ b/src/paimon/common/reader/blob_view_resolving_batch_reader.cpp @@ -41,9 +41,8 @@ namespace paimon { BlobViewResolvingBatchReader::BlobViewResolvingBatchReader( std::unique_ptr&& reader, std::vector read_blob_view_fields, - BlobViewResolver resolver, const std::shared_ptr& pool) - : pool_(pool), - arrow_pool_(GetArrowPool(pool)), + BlobViewResolver resolver, const std::shared_ptr& arrow_pool) + : arrow_pool_(arrow_pool), reader_(std::move(reader)), read_blob_view_fields_(std::make_move_iterator(read_blob_view_fields.begin()), std::make_move_iterator(read_blob_view_fields.end())), @@ -91,6 +90,7 @@ Result BlobViewResolvingBatchReader::NextBatch() { arrow::StructArray::Make(new_fields, field_names)); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*resolved_struct_array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return batch; } diff --git a/src/paimon/common/reader/blob_view_resolving_batch_reader.h b/src/paimon/common/reader/blob_view_resolving_batch_reader.h index e3b14abb2..b97131a82 100644 --- a/src/paimon/common/reader/blob_view_resolving_batch_reader.h +++ b/src/paimon/common/reader/blob_view_resolving_batch_reader.h @@ -38,7 +38,7 @@ class BlobViewResolvingBatchReader : public BatchReader { BlobViewResolvingBatchReader(std::unique_ptr&& reader, std::vector read_blob_view_fields, BlobViewResolver resolver, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result NextBatch() override; @@ -55,8 +55,7 @@ class BlobViewResolvingBatchReader : public BatchReader { const std::shared_ptr& blob_view_struct_array); private: - std::shared_ptr pool_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::unique_ptr reader_; std::set read_blob_view_fields_; BlobViewResolver resolver_; diff --git a/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp b/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp index 4091718fc..bdc4bbc6a 100644 --- a/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp +++ b/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp @@ -37,6 +37,7 @@ #include "paimon/common/data/blob_utils.h" #include "paimon/common/data/blob_view_struct.h" #include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/memory/bytes.h" #include "paimon/memory/memory_pool.h" @@ -52,6 +53,7 @@ class BlobViewResolvingBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } void TearDown() override { @@ -130,6 +132,7 @@ class BlobViewResolvingBatchReaderTest : public ::testing::Test { private: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; TEST_F(BlobViewResolvingBatchReaderTest, TestEofBatch) { @@ -138,7 +141,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestEofBatch) { return std::shared_ptr(); }); BlobViewResolvingBatchReader reader(std::move(inner_reader), {"blob_col"}, std::move(resolver), - pool_); + arrow_pool_); ASSERT_OK_AND_ASSIGN(auto batch, reader.NextBatch()); ASSERT_TRUE(BatchReader::IsEofBatch(batch)); } @@ -155,9 +158,10 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestEmptyReadBlobViewFields) { }); auto inner_reader = std::make_unique(struct_array); - BlobViewResolvingBatchReader reader(std::move(inner_reader), /*read_blob_view_fields=*/{}, - std::move(resolver), pool_); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(&reader)); + auto reader = std::make_unique( + std::move(inner_reader), /*read_blob_view_fields=*/std::vector(), + std::move(resolver), arrow_pool_); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto expected_array = std::make_shared(struct_array); ASSERT_TRUE(expected_array->Equals(*result_array)); ASSERT_FALSE(resolver_called); @@ -186,9 +190,10 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestResolvesBlobViewColumn) { }); auto inner_reader = std::make_unique(src_struct); - BlobViewResolvingBatchReader reader(std::move(inner_reader), {"blob_col"}, std::move(resolver), - pool_); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(&reader)); + auto reader = std::make_unique( + std::move(inner_reader), std::vector{"blob_col"}, std::move(resolver), + arrow_pool_); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto struct_array = std::dynamic_pointer_cast(result_array->chunk(0)); auto result_blob_column = @@ -207,7 +212,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestResolverError) { }); auto inner_reader = std::make_unique(src_struct); BlobViewResolvingBatchReader reader(std::move(inner_reader), {"blob_col"}, std::move(resolver), - pool_); + arrow_pool_); ASSERT_NOK_WITH_MSG(reader.NextBatch(), "cache miss"); } diff --git a/src/paimon/common/reader/complete_row_kind_batch_reader.cpp b/src/paimon/common/reader/complete_row_kind_batch_reader.cpp index b613514e6..f7d8e2ac2 100644 --- a/src/paimon/common/reader/complete_row_kind_batch_reader.cpp +++ b/src/paimon/common/reader/complete_row_kind_batch_reader.cpp @@ -32,6 +32,7 @@ #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/status.h" @@ -41,7 +42,10 @@ namespace paimon { Result CompleteRowKindBatchReader::NextBatch() { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, NextBatchWithBitmap()); - return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool_.get()); + PAIMON_ASSIGN_OR_RAISE( + BatchReader::ReadBatch batch, + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool_)); + return batch; } Result CompleteRowKindBatchReader::NextBatchWithBitmap() { @@ -77,6 +81,7 @@ Result CompleteRowKindBatchReader::NextBatchWi arrow::StructArray::Make(fields_with_row_kind, field_names_with_row_kind_)); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*array_with_row_kind, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return batch_with_bitmap; } diff --git a/src/paimon/common/reader/complete_row_kind_batch_reader.h b/src/paimon/common/reader/complete_row_kind_batch_reader.h index 021dfb4a3..e2581aba3 100644 --- a/src/paimon/common/reader/complete_row_kind_batch_reader.h +++ b/src/paimon/common/reader/complete_row_kind_batch_reader.h @@ -42,8 +42,8 @@ class Metrics; class CompleteRowKindBatchReader : public BatchReader { public: CompleteRowKindBatchReader(std::unique_ptr&& reader, - const std::shared_ptr& pool) - : arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)) {} + const std::shared_ptr& arrow_pool) + : arrow_pool_(arrow_pool), reader_(std::move(reader)) {} Result NextBatch() override; @@ -65,7 +65,7 @@ class CompleteRowKindBatchReader : public BatchReader { void UpdateFieldNamesWithRowKind(const std::shared_ptr& struct_array); private: - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::unique_ptr reader_; std::shared_ptr row_kind_array_; std::vector field_names_with_row_kind_; diff --git a/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp b/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp index ab0dc2a78..c12830102 100644 --- a/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp +++ b/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp @@ -30,6 +30,7 @@ #include "gtest/gtest.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/file_format.h" #include "paimon/format/file_format_factory.h" #include "paimon/fs/file_system.h" @@ -46,6 +47,7 @@ class CompleteRowKindBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } void TearDown() override { @@ -67,18 +69,21 @@ class CompleteRowKindBatchReaderTest : public ::testing::Test { EXPECT_TRUE(arrow_status.ok()); EXPECT_OK(orc_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - return std::make_unique(std::move(orc_batch_reader), pool_); + return std::make_unique(std::move(orc_batch_reader), + arrow_pool_); } std::unique_ptr PrepareCompleteRowKindBatchReader( const std::shared_ptr& src_array, int32_t batch_size) const { auto file_batch_reader = std::make_unique(src_array, src_array->type(), batch_size); - return std::make_unique(std::move(file_batch_reader), pool_); + return std::make_unique(std::move(file_batch_reader), + arrow_pool_); } private: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; TEST_F(CompleteRowKindBatchReaderTest, TestSimple) { @@ -92,7 +97,7 @@ TEST_F(CompleteRowKindBatchReaderTest, TestSimple) { auto read_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); auto reader = PrepareCompleteRowKindBatchReader(file_name, read_schema, /*batch_size=*/1); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected_array; std::vector result_fields = read_fields; @@ -106,7 +111,6 @@ TEST_F(CompleteRowKindBatchReaderTest, TestSimple) { &expected_array); ASSERT_TRUE(array_status.ok()); ASSERT_TRUE(expected_array->Equals(*result_array)); - reader->Close(); } TEST_F(CompleteRowKindBatchReaderTest, TestInnerReaderContainsRowKind) { @@ -121,7 +125,7 @@ TEST_F(CompleteRowKindBatchReaderTest, TestInnerReaderContainsRowKind) { auto read_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); auto reader = PrepareCompleteRowKindBatchReader(file_name, read_schema, /*batch_size=*/1); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected_array; auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -132,7 +136,6 @@ TEST_F(CompleteRowKindBatchReaderTest, TestInnerReaderContainsRowKind) { &expected_array); ASSERT_TRUE(array_status.ok()); ASSERT_TRUE(expected_array->Equals(*result_array)); - reader->Close(); } TEST_F(CompleteRowKindBatchReaderTest, TestNestedType) { @@ -151,7 +154,7 @@ TEST_F(CompleteRowKindBatchReaderTest, TestNestedType) { .ValueOrDie(); ASSERT_TRUE(src_array); auto reader = PrepareCompleteRowKindBatchReader(src_array, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected_array; arrow::FieldVector read_fields = fields; @@ -166,7 +169,6 @@ TEST_F(CompleteRowKindBatchReaderTest, TestNestedType) { &expected_array); ASSERT_TRUE(array_status.ok()); ASSERT_TRUE(expected_array->Equals(*result_array)); - reader->Close(); } } // namespace paimon::test diff --git a/src/paimon/common/reader/concat_batch_reader.cpp b/src/paimon/common/reader/concat_batch_reader.cpp index c36886c51..9129d5702 100644 --- a/src/paimon/common/reader/concat_batch_reader.cpp +++ b/src/paimon/common/reader/concat_batch_reader.cpp @@ -30,13 +30,16 @@ namespace paimon { class MemoryPool; ConcatBatchReader::ConcatBatchReader(std::vector>&& readers, - const std::shared_ptr& pool) - : arrow_pool_(GetArrowPool(pool)), readers_(std::move(readers)), current_(0) {} + const std::shared_ptr& arrow_pool) + : arrow_pool_(arrow_pool), readers_(std::move(readers)), current_(0) {} Result ConcatBatchReader::NextBatch() { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, NextBatchWithBitmap()); - return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool_.get()); + PAIMON_ASSIGN_OR_RAISE( + BatchReader::ReadBatch batch, + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool_)); + return batch; } void ConcatBatchReader::Close() { diff --git a/src/paimon/common/reader/concat_batch_reader.h b/src/paimon/common/reader/concat_batch_reader.h index 7a1b6a19c..4ce308aa3 100644 --- a/src/paimon/common/reader/concat_batch_reader.h +++ b/src/paimon/common/reader/concat_batch_reader.h @@ -36,7 +36,7 @@ class MemoryPool; class ConcatBatchReader : public BatchReader { public: ConcatBatchReader(std::vector>&& readers, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result NextBatch() override; Result NextBatchWithBitmap() override; @@ -44,7 +44,7 @@ class ConcatBatchReader : public BatchReader { std::shared_ptr GetReaderMetrics() const override; private: - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::vector> readers_; size_t current_; }; diff --git a/src/paimon/common/reader/concat_batch_reader_test.cpp b/src/paimon/common/reader/concat_batch_reader_test.cpp index 37a415dea..85f716a75 100644 --- a/src/paimon/common/reader/concat_batch_reader_test.cpp +++ b/src/paimon/common/reader/concat_batch_reader_test.cpp @@ -29,6 +29,7 @@ #include "arrow/array/array_nested.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/mock/mock_file_batch_reader.h" @@ -40,6 +41,7 @@ namespace paimon::test { class ConcatBatchReaderTest : public ::testing::Test { void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } void CheckResult(const std::vector& batches, const std::string& expected) { std::vector>> batches_with_bitmap; @@ -71,9 +73,10 @@ class ConcatBatchReaderTest : public ::testing::Test { data, data->type(), RoaringBitmap32::From(bitmap_data), batch_size); readers.push_back(std::move(reader)); } - auto concat_reader = std::make_unique(std::move(readers), pool_); + auto concat_reader = + std::make_unique(std::move(readers), arrow_pool_); ASSERT_OK_AND_ASSIGN(auto result_chunk_array, - ReadResultCollector::CollectResult(concat_reader.get())); + ReadResultCollector::CollectResult(std::move(concat_reader))); if (expected.empty()) { ASSERT_FALSE(result_chunk_array); return; @@ -91,6 +94,7 @@ class ConcatBatchReaderTest : public ::testing::Test { private: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; TEST_F(ConcatBatchReaderTest, TestSimple) { CheckResult({"[10, 11, 12, 13, 14]"}, "[10, 11, 12, 13, 14]"); diff --git a/src/paimon/common/reader/data_evolution_file_reader.cpp b/src/paimon/common/reader/data_evolution_file_reader.cpp index 6f2e12d4a..55e3aabb9 100644 --- a/src/paimon/common/reader/data_evolution_file_reader.cpp +++ b/src/paimon/common/reader/data_evolution_file_reader.cpp @@ -36,7 +36,7 @@ Result> DataEvolutionFileReader::Create std::vector>&& readers, const std::shared_ptr& read_schema, int32_t read_batch_size, const std::vector& reader_offsets, const std::vector& field_offsets, - const std::shared_ptr& pool) { + const std::shared_ptr& arrow_pool) { if (read_schema->num_fields() == 0) { return Status::Invalid("read schema must not be empty"); } @@ -55,7 +55,7 @@ Result> DataEvolutionFileReader::Create } return std::unique_ptr( new DataEvolutionFileReader(std::move(readers), read_schema, read_batch_size, - reader_offsets, field_offsets, GetArrowPool(pool))); + reader_offsets, field_offsets, arrow_pool)); } Result DataEvolutionFileReader::NextBatchWithBitmap() { @@ -103,6 +103,7 @@ Result DataEvolutionFileReader::NextBatchWithB std::unique_ptr target_c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*target_array, target_c_arrow_array.get(), target_c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(target_c_arrow_array.get(), arrow_pool_)); auto target_batch = std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema)); return ReaderUtils::AddAllValidBitmap(std::move(target_batch)); } diff --git a/src/paimon/common/reader/data_evolution_file_reader.h b/src/paimon/common/reader/data_evolution_file_reader.h index 76494358c..31743d1ef 100644 --- a/src/paimon/common/reader/data_evolution_file_reader.h +++ b/src/paimon/common/reader/data_evolution_file_reader.h @@ -59,7 +59,7 @@ class DataEvolutionFileReader : public BatchReader { std::vector>&& readers, const std::shared_ptr& read_schema, int32_t read_batch_size, const std::vector& reader_offsets, const std::vector& field_offsets, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result NextBatch() override { return Status::Invalid( diff --git a/src/paimon/common/reader/data_evolution_file_reader_test.cpp b/src/paimon/common/reader/data_evolution_file_reader_test.cpp index 9e545c544..bae51a3f8 100644 --- a/src/paimon/common/reader/data_evolution_file_reader_test.cpp +++ b/src/paimon/common/reader/data_evolution_file_reader_test.cpp @@ -40,6 +40,7 @@ class DataEvolutionFileReaderTest : public ::testing::Test, public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } void TearDown() override { @@ -77,7 +78,7 @@ class DataEvolutionFileReaderTest : public ::testing::Test, ASSERT_OK_AND_ASSIGN( auto data_evolution_file_reader, DataEvolutionFileReader::Create(std::move(readers), read_schema, batch_size, - reader_offsets, field_offsets, pool_)); + reader_offsets, field_offsets, arrow_pool_)); // check metrics, data_evolution_file_reader collects all row of each // MockFileBatchReader auto metrics = data_evolution_file_reader->GetReaderMetrics(); @@ -85,10 +86,9 @@ class DataEvolutionFileReaderTest : public ::testing::Test, "{\"mock.number.of.rows\":" + std::to_string(total_row_count) + "}"); // check result array - ASSERT_OK_AND_ASSIGN( - auto result_array, - paimon::test::ReadResultCollector::CollectResult(data_evolution_file_reader.get())); - data_evolution_file_reader->Close(); + ASSERT_OK_AND_ASSIGN(auto result_array, + paimon::test::ReadResultCollector::CollectResult( + std::move(data_evolution_file_reader))); auto expected_chunk_array = std::make_shared(expected_array); ASSERT_TRUE(result_array->Equals(expected_chunk_array)); } @@ -112,7 +112,7 @@ class DataEvolutionFileReaderTest : public ::testing::Test, readers.push_back(std::move(file_batch_reader)); DataEvolutionFileReader fake_data_evolution_reader( std::move(readers), /*read_schema=*/arrow::schema({}), read_batch_size, - /*reader_offsets=*/{}, /*field_offsets=*/{}, GetArrowPool(pool_)); + /*reader_offsets=*/{}, /*field_offsets=*/{}, arrow_pool_); arrow::ArrayVector result_array_vec; while (true) { ASSERT_OK_AND_ASSIGN(auto result_array, @@ -142,15 +142,16 @@ class DataEvolutionFileReaderTest : public ::testing::Test, private: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; TEST_F(DataEvolutionFileReaderTest, TestInvalid) { { arrow::FieldVector read_fields; auto read_schema = arrow::schema(read_fields); - ASSERT_NOK_WITH_MSG( - DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, {}, {}, pool_), - "read schema must not be empty"); + ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, + {}, {}, arrow_pool_), + "read schema must not be empty"); } { arrow::FieldVector read_fields = { @@ -162,9 +163,10 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { auto read_schema = arrow::schema(read_fields); std::vector reader_offsets = {0, 0, 1}; std::vector field_offsets = {0, 1, 0}; - ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, - reader_offsets, field_offsets, pool_), - "read schema, row offsets and field offsets must have the same size"); + ASSERT_NOK_WITH_MSG( + DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, reader_offsets, + field_offsets, arrow_pool_), + "read schema, row offsets and field offsets must have the same size"); } { arrow::FieldVector read_fields = { @@ -176,9 +178,10 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { auto read_schema = arrow::schema(read_fields); std::vector reader_offsets = {0, 0, 1, 1}; std::vector field_offsets = {0, 1, 1, 0}; - ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, - reader_offsets, field_offsets, pool_), - "readers must not be empty"); + ASSERT_NOK_WITH_MSG( + DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, reader_offsets, + field_offsets, arrow_pool_), + "readers must not be empty"); } { std::vector> readers; @@ -195,7 +198,7 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { std::vector field_offsets = {0, 1, 1, 0}; ASSERT_NOK_WITH_MSG( DataEvolutionFileReader::Create(std::move(readers), read_schema, /*read_batch_size=*/10, - reader_offsets, field_offsets, pool_), + reader_offsets, field_offsets, arrow_pool_), "reader offset is out of range of readers"); } } @@ -637,10 +640,10 @@ TEST_P(DataEvolutionFileReaderTest, TestSingleReaderRowCountMismatch) { ASSERT_OK_AND_ASSIGN( auto data_evolution_file_reader, DataEvolutionFileReader::Create(std::move(readers), read_schema, /*read_batch_size=*/10, - reader_offsets, field_offsets, pool_)); + reader_offsets, field_offsets, arrow_pool_)); // array0 has 6 rows but array1 only has 5 rows ASSERT_NOK_WITH_MSG( - paimon::test::ReadResultCollector::CollectResult(data_evolution_file_reader.get()), + paimon::test::ReadResultCollector::CollectResult(std::move(data_evolution_file_reader)), "array for single reader length mismatch others"); } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index ed7fa304d..6fab25e04 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -45,20 +45,16 @@ namespace paimon { Result> LateMaterializingFileBatchReader::Create( - std::unique_ptr inner, std::shared_ptr pool) { - // The reader's own compaction allocations go through an arrow pool; bridge the paimon pool - // once here so the accounting matches the rest of the read path. - if (pool == nullptr) { - return Status::Invalid("pool could not be nullptr."); + std::unique_ptr inner, const std::shared_ptr& arrow_pool) { + if (arrow_pool == nullptr) { + return Status::Invalid("arrow pool could not be nullptr."); } if (inner == nullptr) { return Status::Invalid("inner could not be nullptr."); } auto* prefetch_inner = dynamic_cast(inner.get()); - std::shared_ptr arrow_pool = GetArrowPool(pool); - auto reader = - std::unique_ptr(new LateMaterializingFileBatchReader( - std::move(inner), prefetch_inner, std::move(arrow_pool))); + auto reader = std::unique_ptr( + new LateMaterializingFileBatchReader(std::move(inner), prefetch_inner, arrow_pool)); return reader; } @@ -241,6 +237,7 @@ Result LateMaterializingFileBatchReader::AssembleFul std::unique_ptr<::ArrowSchema> c_schema = std::make_unique<::ArrowSchema>(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*full_struct, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return std::make_pair(std::move(c_array), std::move(c_schema)); } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 231625db6..7a225bca7 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -42,7 +42,8 @@ class PredicateFilter; class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { public: static Result> Create( - std::unique_ptr inner, std::shared_ptr pool); + std::unique_ptr inner, + const std::shared_ptr& arrow_pool); Result NextBatch() override; diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index 075fbe107..e75870157 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -35,6 +35,7 @@ #include "paimon/common/reader/late_materializing_reader_builder.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" #include "paimon/common/reader/reader_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/read_ahead_cache.h" @@ -58,6 +59,8 @@ namespace paimon::test { class LateMaterializingFileBatchReaderTest : public ::testing::Test { public: void SetUp() override { + pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); k_field_ = arrow::field("k", arrow::int64()); v_field_ = arrow::field("v", arrow::utf8()); full_fields_ = {k_field_, v_field_}; @@ -142,6 +145,18 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { return arrow::internal::checked_pointer_cast(combined); } + Result> CollectStruct( + std::unique_ptr reader) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr chunked, + ReadResultCollector::CollectResult(std::move(reader))); + if (chunked == nullptr) { + return std::shared_ptr(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr combined, + arrow::Concatenate(chunked->chunks())); + return arrow::internal::checked_pointer_cast(combined); + } + // Build a struct with 5 columns [a:int64, b:utf8, c:int64, d:utf8, e:int64], each carrying a // distinct value pattern so any column reordering is detected. std::shared_ptr BuildMultiFieldData(int32_t n) { @@ -201,6 +216,8 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { } protected: + std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::shared_ptr k_field_; std::shared_ptr v_field_; arrow::FieldVector full_fields_; @@ -211,8 +228,8 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr, std::nullopt)); @@ -230,8 +247,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); // read schema is just {k}; the predicate on k covers all columns -> payload empty auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(0l)); @@ -251,8 +268,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { TEST_F(LateMaterializingFileBatchReaderTest, ContiguousSubsetAcrossBatches) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(4l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -276,8 +293,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ScatteredAlternatingMatch) { } auto data = BuildData(ks); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -300,8 +317,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { } auto data = BuildData(ks); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); // predicate hits {1,3,5,7,9,11}; selection keeps only {1,5,9} @@ -326,8 +343,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(100l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -340,8 +357,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { TEST_F(LateMaterializingFileBatchReaderTest, SeekToRowRealignsProbeCursor) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(5l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -369,8 +386,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); auto* mock_ptr = mock.get(); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(2l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -398,8 +415,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); auto predicate1 = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(7l)); @@ -429,8 +446,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { TEST_F(LateMaterializingFileBatchReaderTest, ForwardsRowCountAndFileSchema) { auto data = BuildData({0, 1, 2, 3}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); ASSERT_OK_AND_ASSIGN(uint64_t num_rows, reader->GetNumberOfRows()); EXPECT_EQ(num_rows, 4u); @@ -446,8 +463,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { auto data = BuildMultiFieldData(10); auto type = data->type(); auto mock = std::make_unique(data, type, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); // probe columns = {a (idx0), c (idx2)}; payload columns = {b, d, e} auto pred_a = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, Literal(3l)); @@ -456,7 +473,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({pred_a, pred_c})); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(reader.get())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + CollectStruct(std::move(reader))); ASSERT_TRUE(result); // a >= 3 and c(=i*100) < 700 -> i in {3,4,5,6} ASSERT_EQ(result->length(), 4); @@ -489,14 +507,15 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { auto data = BuildNestedData(8); auto type = data->type(); auto mock = std::make_unique(data, type, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); // probe = {k}; payload = {arr (list), tag} auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(reader.get())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + CollectStruct(std::move(reader))); ASSERT_TRUE(result); ASSERT_EQ(result->length(), 3); // k = 5,6,7 auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); @@ -521,8 +540,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); LateMaterializingReaderBuilder builder( - std::make_unique(data, full_type_, /*batch_size=*/3), - GetDefaultPool()); + std::make_unique(data, full_type_, /*batch_size=*/3), arrow_pool_); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); ASSERT_OK_AND_ASSIGN( @@ -532,14 +550,15 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, GetDefaultPool())); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l)); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(impl.get())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + CollectStruct(std::move(impl))); ASSERT_TRUE(result); ASSERT_EQ(result->length(), 6); // k = 4..9 auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); @@ -549,7 +568,6 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { EXPECT_EQ(k->Value(j), 4 + j); EXPECT_EQ(v->GetString(j), "v_" + std::to_string(4 + j)); } - impl->Close(); } // Re-setting the read schema on the prefetch impl (which re-broadcasts to the inner LM readers and @@ -557,8 +575,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); LateMaterializingReaderBuilder builder( - std::make_unique(data, full_type_, /*batch_size=*/3), - GetDefaultPool()); + std::make_unique(data, full_type_, /*batch_size=*/3), arrow_pool_); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); ASSERT_OK_AND_ASSIGN( @@ -568,7 +585,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, GetDefaultPool())); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto full_schema = arrow::schema(full_fields_); auto predicate1 = @@ -620,8 +637,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee // Per-batch ranges (the mock's default) let the impl split work across the parallel readers, // and each range-honoring reader only reads its assigned slice. LateMaterializingReaderBuilder builder( - std::make_unique(data, full_type_, /*batch_size=*/3), - GetDefaultPool()); + std::make_unique(data, full_type_, /*batch_size=*/3), arrow_pool_); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(3)); ASSERT_OK_AND_ASSIGN( @@ -631,14 +647,15 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee /*prefetch_max_parallel_num=*/3, /*batch_size=*/3, /*prefetch_batch_count=*/6, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, GetDefaultPool())); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(impl.get())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + CollectStruct(std::move(impl))); ASSERT_TRUE(result); ASSERT_EQ(result->length(), 15); // k = 5..19 auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); @@ -648,7 +665,6 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee EXPECT_EQ(k->Value(j), 5 + j); EXPECT_EQ(v->GetString(j), "v_" + std::to_string(5 + j)); } - impl->Close(); } // When the predicate's field type does not match the probe schema, the @@ -657,8 +673,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee TEST_F(LateMaterializingFileBatchReaderTest, FailsOnPredicateTypeMismatch) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN( - auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); // k is int64 in the schema, but the predicate claims FieldType::INT (int32). auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::INT, Literal(10)); diff --git a/src/paimon/common/reader/late_materializing_reader_builder.h b/src/paimon/common/reader/late_materializing_reader_builder.h index 5c7be5077..56571fcb9 100644 --- a/src/paimon/common/reader/late_materializing_reader_builder.h +++ b/src/paimon/common/reader/late_materializing_reader_builder.h @@ -33,11 +33,10 @@ namespace paimon { class LateMaterializingReaderBuilder : public ReaderBuilder { public: LateMaterializingReaderBuilder(std::unique_ptr inner, - std::shared_ptr pool) - : inner_(std::move(inner)), pool_(std::move(pool)) {} + std::shared_ptr arrow_pool) + : inner_(std::move(inner)), arrow_pool_(std::move(arrow_pool)) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { - pool_ = pool; inner_->WithMemoryPool(pool); return this; } @@ -58,13 +57,13 @@ class LateMaterializingReaderBuilder : public ReaderBuilder { inner_->Build(stream)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr reader, - LateMaterializingFileBatchReader::Create(std::move(format_reader), pool_)); + LateMaterializingFileBatchReader::Create(std::move(format_reader), arrow_pool_)); return std::unique_ptr(std::move(reader)); } private: std::unique_ptr inner_; - std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; } // namespace paimon diff --git a/src/paimon/common/reader/predicate_batch_reader.cpp b/src/paimon/common/reader/predicate_batch_reader.cpp index b4d6b91d0..621d7568c 100644 --- a/src/paimon/common/reader/predicate_batch_reader.cpp +++ b/src/paimon/common/reader/predicate_batch_reader.cpp @@ -48,12 +48,12 @@ class MemoryPool; PredicateBatchReader::PredicateBatchReader(std::unique_ptr&& reader, const std::shared_ptr& predicate, - const std::shared_ptr& pool) - : arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)), predicate_(predicate) {} + const std::shared_ptr& arrow_pool) + : arrow_pool_(arrow_pool), reader_(std::move(reader)), predicate_(predicate) {} Result> PredicateBatchReader::Create( std::unique_ptr&& reader, const std::shared_ptr& predicate, - const std::shared_ptr& pool) { + const std::shared_ptr& arrow_pool) { if (!predicate) { return Status::Invalid("create predicate batch reader failed. predicate is nullptr"); } @@ -62,13 +62,16 @@ Result> PredicateBatchReader::Create( fmt::format("predicate {} does not support Test", predicate->ToString())); } return std::unique_ptr( - new PredicateBatchReader(std::move(reader), predicate, pool)); + new PredicateBatchReader(std::move(reader), predicate, arrow_pool)); } Result PredicateBatchReader::NextBatch() { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, NextBatchWithBitmap()); - return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool_.get()); + PAIMON_ASSIGN_OR_RAISE( + BatchReader::ReadBatch batch, + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool_)); + return batch; } Result PredicateBatchReader::NextBatchWithBitmap() { diff --git a/src/paimon/common/reader/predicate_batch_reader.h b/src/paimon/common/reader/predicate_batch_reader.h index 0211a7d1e..65579f5bd 100644 --- a/src/paimon/common/reader/predicate_batch_reader.h +++ b/src/paimon/common/reader/predicate_batch_reader.h @@ -40,7 +40,7 @@ class PredicateBatchReader : public BatchReader { public: static Result> Create( std::unique_ptr&& reader, const std::shared_ptr& predicate, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); ~PredicateBatchReader() override = default; @@ -59,12 +59,12 @@ class PredicateBatchReader : public BatchReader { private: PredicateBatchReader(std::unique_ptr&& reader, const std::shared_ptr& predicate, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Status BindPredicateToArray(const arrow::Array& array); Result Filter(const std::shared_ptr& array); private: - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::unique_ptr reader_; std::shared_ptr predicate_; std::shared_ptr predicate_filter_; diff --git a/src/paimon/common/reader/predicate_batch_reader_test.cpp b/src/paimon/common/reader/predicate_batch_reader_test.cpp index 8bcc53af8..3b6c70c9a 100644 --- a/src/paimon/common/reader/predicate_batch_reader_test.cpp +++ b/src/paimon/common/reader/predicate_batch_reader_test.cpp @@ -31,6 +31,7 @@ #include "arrow/array/builder_primitive.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/defs.h" #include "paimon/memory/memory_pool.h" @@ -78,11 +79,11 @@ class PredicateBatchReaderTest : public ::testing::Test { void CheckResult(std::unique_ptr&& reader, const std::shared_ptr& predicate, const std::shared_ptr& expected_array) const { - ASSERT_OK_AND_ASSIGN( - auto predicate_reader, - PredicateBatchReader::Create(std::move(reader), predicate, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto predicate_reader, + PredicateBatchReader::Create(std::move(reader), predicate, + GetSharedArrowPool(GetDefaultPool()))); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(predicate_reader.get())); + ReadResultCollector::CollectResult(std::move(predicate_reader))); if (expected_array) { ASSERT_TRUE(result_array->Equals(expected_array)); } else { @@ -204,7 +205,8 @@ TEST_F(PredicateBatchReaderTest, TestFullAndEmptyCase) { TEST_F(PredicateBatchReaderTest, TestInvalidInput) { auto data_array = PrepareArray(8); auto reader = std::make_unique(data_array, data_type_, /*batch_size=*/10); - ASSERT_NOK_WITH_MSG(PredicateBatchReader::Create(std::move(reader), nullptr, GetDefaultPool()), + ASSERT_NOK_WITH_MSG(PredicateBatchReader::Create(std::move(reader), nullptr, + GetSharedArrowPool(GetDefaultPool())), "create predicate batch reader failed. predicate is nullptr"); } diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp index 5285a2d80..859ceddd7 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -205,7 +205,8 @@ Result> PrefetchFileBatchReaderImpl uint32_t prefetch_batch_count, bool enable_adaptive_prefetch_strategy, const std::shared_ptr& executor, bool initialize_read_ranges, bool read_ahead_cache_enabled, const CacheConfig& cache_config, bool enable_io_metrics, - const std::shared_ptr& pool) { + const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool) { if (prefetch_max_parallel_num == 0) { return Status::Invalid("prefetch max parallel num should be greater than 0."); } @@ -276,7 +277,7 @@ Result> PrefetchFileBatchReaderImpl auto reader = std::unique_ptr(new PrefetchFileBatchReaderImpl( readers, batch_size, prefetch_queue_capacity, enable_adaptive_prefetch_strategy, executor, - cache, io_metrics, pool)); + cache, io_metrics, arrow_pool)); if (initialize_read_ranges) { // normally initialize read ranges should be false, as set read schema will refresh read // ranges, and set read schema will always be called before read. @@ -290,12 +291,12 @@ PrefetchFileBatchReaderImpl::PrefetchFileBatchReaderImpl( uint32_t prefetch_queue_capacity, bool enable_adaptive_prefetch_strategy, const std::shared_ptr& executor, const std::shared_ptr& cache, const std::shared_ptr& io_metrics, - const std::shared_ptr& pool) + const std::shared_ptr& arrow_pool) : readers_(std::move(readers)), batch_size_(batch_size), executor_(executor), cache_(cache), - arrow_pool_(GetArrowPool(pool)), + arrow_pool_(arrow_pool), prefetch_queue_capacity_(prefetch_queue_capacity), enable_adaptive_prefetch_strategy_(enable_adaptive_prefetch_strategy), prefetch_metrics_(std::make_shared()), @@ -654,6 +655,7 @@ Status PrefetchFileBatchReaderImpl::HandleReadResult( ArrowUtils::NormalizeArrayOffsets(sliced_array, arrow_pool_.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*normalized_array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); bitmap.RemoveRange(slice_end, array->length()); global_row_ids = std::vector(global_row_ids.begin(), global_row_ids.begin() + slice_end); diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h index c99323698..37bf2f896 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h @@ -68,7 +68,8 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { uint32_t prefetch_max_parallel_num, int32_t batch_size, uint32_t prefetch_batch_count, bool enable_adaptive_prefetch_strategy, const std::shared_ptr& executor, bool initialize_read_ranges, bool read_ahead_cache_enabled, const CacheConfig& cache_config, - bool enable_io_metrics, const std::shared_ptr& pool); + bool enable_io_metrics, const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool); ~PrefetchFileBatchReaderImpl() override; @@ -122,7 +123,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { uint32_t prefetch_queue_capacity, bool enable_adaptive_prefetch_strategy, const std::shared_ptr& executor, const std::shared_ptr& cache, const std::shared_ptr& io_metrics, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Status CleanUp(); void Workloop(); @@ -169,7 +170,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { std::condition_variable cv_; std::shared_ptr executor_; std::shared_ptr cache_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; mutable std::shared_mutex rw_mutex_; std::unique_ptr background_thread_; diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp index 514e7e49d..7e95e9f4d 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp @@ -25,6 +25,7 @@ #include "arrow/compute/api.h" #include "arrow/ipc/api.h" #include "gtest/gtest.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/read_ahead_cache.h" @@ -179,6 +180,8 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, public ::testing::WithParamInterface { public: void SetUp() override { + pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); fields_ = {arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int64()), arrow::field("f2", arrow::boolean())}; data_type_ = arrow::struct_(fields_); @@ -267,7 +270,7 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, read_ahead_cache_enabled, CacheConfig(), - /*enable_io_metrics=*/true, GetDefaultPool())); + /*enable_io_metrics=*/true, pool_, arrow_pool_)); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); EXPECT_TRUE(arrow_status.ok()); @@ -302,6 +305,8 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, } private: + std::shared_ptr pool_; + std::shared_ptr arrow_pool_; arrow::FieldVector fields_; std::shared_ptr data_type_; std::shared_ptr mock_fs_; @@ -356,7 +361,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestSimple) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); if (prefetch_max_parallel_num == 1) { ASSERT_NOK( reader->GetReaderMetrics()->GetCounter(PrefetchIoMetrics::READ_LATENCY_COUNT)); @@ -383,7 +388,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLimits) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/true, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/true, pool_, arrow_pool_)); // simulate read limits, only read 8 batches for (int32_t i = 0; i < 8; i++) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, @@ -436,7 +441,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithoutInitializeReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); // simulate read limits, only read 8 batches ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "prefetch reader read ranges are not initialized"); @@ -454,7 +459,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestFailedIoMetrics) { /*prefetch_max_parallel_num=*/1, /*batch_size=*/10, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/true, GetDefaultPool())); + /*enable_io_metrics=*/true, pool_, arrow_pool_)); ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "injected synchronous read failure"); std::shared_ptr metrics = reader->GetReaderMetrics(); @@ -542,7 +547,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_OK(prefetch_reader->RefreshReadRanges()); std::vector> read_ranges_0 = {{0, 30}, {90, 101}}; @@ -572,7 +577,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRangesDisablePrefetchByAdapti /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/true, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); ASSERT_FALSE(reader->NeedPrefetch()); std::shared_ptr metrics = reader->GetReaderMetrics(); @@ -596,7 +601,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_FALSE(prefetch_reader->need_prefetch_); prefetch_reader->need_prefetch_ = true; @@ -639,7 +644,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRangesReturnErrorWhenPushDownFail prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->need_prefetch_ = true; @@ -659,13 +664,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, WorkloopSetReadStatusWhenCacheInitFailed /*hole_size_limit=*/8 * 1024, /*pre_buffer_limit=*/128 * 1024); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - invalid_cache_config, /*enable_io_metrics=*/false, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN( + auto reader, PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + invalid_cache_config, /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->Workloop(); @@ -685,7 +690,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenShutdown) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->is_shutdown_ = true; @@ -703,7 +708,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenNoCurrentReadRang prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->read_ranges_in_group_ = {{}}; @@ -721,7 +726,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLargeBatchSize) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -741,7 +746,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPartialReaderSuccessRead) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { dynamic_cast(prefetch_reader->readers_[i].get()) @@ -786,7 +791,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { @@ -824,7 +829,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithEmptyData) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -843,7 +848,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCallNextBatchAfterReadingEof) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -886,7 +891,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCreateReaderWithoutNextBatch) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); } TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { @@ -901,7 +906,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*prefetch_max_parallel_num=*/0, batch_size, 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, GetDefaultPool())); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -909,7 +914,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, /*batch_size=*/-1, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, GetDefaultPool())); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -917,8 +922,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, /*executor=*/nullptr, /*initialize_read_ranges=*/true, - /*read_ahead_cache_enabled=*/true, CacheConfig(), /*enable_io_metrics=*/false, - GetDefaultPool())); + /*read_ahead_cache_enabled=*/true, CacheConfig(), /*enable_io_metrics=*/false, pool_, + arrow_pool_)); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -926,7 +931,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, GetDefaultPool())); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -934,7 +939,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*fs=*/nullptr, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, GetDefaultPool())); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); } { ASSERT_OK_AND_ASSIGN( @@ -943,7 +948,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, GetDefaultPool())); + CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); ASSERT_NOK_WITH_MSG(reader->SeekToRow(/*row_number=*/101), "not support seek to row for prefetch reader"); } @@ -1047,13 +1052,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithBitmap) { /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, GetDefaultPool())); - ASSERT_OK_AND_ASSIGN(auto result_chunk_array, ReadResultCollector::CollectResult(reader.get())); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto result_chunk_array, + ReadResultCollector::CollectResult(std::move(reader))); ASSERT_OK_AND_ASSIGN(auto data_batch, ReadResultCollector::GetReadBatch(data_array)); - ASSERT_OK_AND_ASSIGN(auto expected_batch, ReaderUtils::ApplyBitmapToReadBatch( - std::make_pair(std::move(data_batch), bitmap), - arrow::default_memory_pool())); + ASSERT_OK_AND_ASSIGN(auto expected_batch, + ReaderUtils::ApplyBitmapToReadBatch( + std::make_pair(std::move(data_batch), bitmap), arrow_pool_)); ASSERT_OK_AND_ASSIGN(auto expected_array, ReadResultCollector::GetArray(std::move(expected_batch))); auto expected_chunk_array = std::make_shared(expected_array); diff --git a/src/paimon/common/reader/reader_utils.cpp b/src/paimon/common/reader/reader_utils.cpp index 7bb28c771..1d7b166de 100644 --- a/src/paimon/common/reader/reader_utils.cpp +++ b/src/paimon/common/reader/reader_utils.cpp @@ -30,6 +30,7 @@ #include "arrow/c/abi.h" #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/status.h" #include "paimon/utils/roaring_bitmap32.h" @@ -69,7 +70,8 @@ void ReaderUtils::ReleaseReadBatch(BatchReader::ReadBatch&& batch) { } Result ReaderUtils::ApplyBitmapToReadBatch( - BatchReader::ReadBatchWithBitmap&& batch_with_bitmap, arrow::MemoryPool* arrow_pool) { + BatchReader::ReadBatchWithBitmap&& batch_with_bitmap, + const std::shared_ptr& arrow_pool) { if (BatchReader::IsEofBatch(batch_with_bitmap)) { return std::move(batch_with_bitmap.first); } @@ -92,12 +94,13 @@ Result ReaderUtils::ApplyBitmapToReadBatch( PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector array_vec, GenerateFilteredArrayVector(arrow_array, bitmap)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, - arrow::Concatenate(array_vec, arrow_pool)); + arrow::Concatenate(array_vec, arrow_pool.get())); assert(result && result->length() > 0); std::unique_ptr result_c_array = std::make_unique(); std::unique_ptr result_c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*result, result_c_array.get(), result_c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(result_c_array.get(), arrow_pool)); return make_pair(std::move(result_c_array), std::move(result_c_schema)); } diff --git a/src/paimon/common/reader/reader_utils.h b/src/paimon/common/reader/reader_utils.h index aa77e404e..f256607ec 100644 --- a/src/paimon/common/reader/reader_utils.h +++ b/src/paimon/common/reader/reader_utils.h @@ -45,7 +45,8 @@ class ReaderUtils { /// @return returned array contains all the valid rows in the input array /// This function may trigger data copy. static Result ApplyBitmapToReadBatch( - BatchReader::ReadBatchWithBitmap&& batch_with_bitmap, arrow::MemoryPool* arrow_pool); + BatchReader::ReadBatchWithBitmap&& batch_with_bitmap, + const std::shared_ptr& arrow_pool); /// @param batch a read batch /// @return return the input batch and a all valid bitmap static BatchReader::ReadBatchWithBitmap AddAllValidBitmap(BatchReader::ReadBatch&& batch); diff --git a/src/paimon/common/reader/reader_utils_test.cpp b/src/paimon/common/reader/reader_utils_test.cpp index 96b5a4da0..ad0ee3b05 100644 --- a/src/paimon/common/reader/reader_utils_test.cpp +++ b/src/paimon/common/reader/reader_utils_test.cpp @@ -30,12 +30,38 @@ #include "arrow/c/abi.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/status.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" #include "paimon/utils/roaring_bitmap32.h" namespace paimon::test { +namespace { + +class ErrorAfterOneBatchReader : public BatchReader { + public: + explicit ErrorAfterOneBatchReader(ReadBatch batch) : batch_(std::move(batch)) {} + + Result NextBatch() override { + if (batch_.first) { + return std::move(batch_); + } + return Status::IOError("expected test error"); + } + + std::shared_ptr GetReaderMetrics() const override { + return nullptr; + } + + void Close() override {} + + private: + ReadBatch batch_; +}; + +} // namespace + TEST(ReaderUtilsTest, TestAddAllValidBitmap) { auto check_result = [](const std::string& src_str) { if (src_str.empty()) { @@ -59,15 +85,32 @@ TEST(ReaderUtilsTest, TestAddAllValidBitmap) { check_result("[10, 20, 30]"); check_result(""); } + +TEST(ReaderUtilsTest, TestCollectResultReleasesBufferedBatchesOnError) { + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[1]").ValueOrDie(); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, ReadResultCollector::GetReadBatch(array)); + + std::shared_ptr lifetime = std::make_shared(1); + std::weak_ptr weak_lifetime = lifetime; + ASSERT_OK(AddArrowArrayLifetime(batch.first.get(), lifetime)); + lifetime.reset(); + + auto reader = std::make_unique(std::move(batch)); + ASSERT_NOK_WITH_MSG(ReadResultCollector::CollectResult(std::move(reader)), + "expected test error"); + ASSERT_TRUE(weak_lifetime.expired()); +} + TEST(ReaderUtilsTest, TestApplyBitmapToReadBatch) { auto check_result = [](const std::string& src_str, const std::vector& bitmap_vec, const std::string& target_str, const std::string& erro_msg = "") { + std::shared_ptr arrow_pool(arrow::default_memory_pool(), + [](arrow::MemoryPool*) {}); auto bitmap = RoaringBitmap32::From(bitmap_vec); if (src_str.empty()) { auto batch_with_bitmap = std::make_pair(BatchReader::MakeEofBatch(), std::move(bitmap)); - ASSERT_OK_AND_ASSIGN(auto result_batch, - ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), - arrow::default_memory_pool())); + ASSERT_OK_AND_ASSIGN(auto result_batch, ReaderUtils::ApplyBitmapToReadBatch( + std::move(batch_with_bitmap), arrow_pool)); ASSERT_TRUE(BatchReader::IsEofBatch(result_batch)); return; } @@ -78,14 +121,13 @@ TEST(ReaderUtilsTest, TestApplyBitmapToReadBatch) { ASSERT_OK_AND_ASSIGN(auto src_batch, ReadResultCollector::GetReadBatch(src_array)); auto batch_with_bitmap = std::make_pair(std::move(src_batch), std::move(bitmap)); if (!erro_msg.empty()) { - ASSERT_NOK_WITH_MSG(ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), - arrow::default_memory_pool()), - erro_msg); + ASSERT_NOK_WITH_MSG( + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool), + erro_msg); return; } - ASSERT_OK_AND_ASSIGN(auto result_batch, - ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), - arrow::default_memory_pool())); + ASSERT_OK_AND_ASSIGN(auto result_batch, ReaderUtils::ApplyBitmapToReadBatch( + std::move(batch_with_bitmap), arrow_pool)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::GetArray(std::move(result_batch))); ASSERT_TRUE(result_array->Equals(target_array)); diff --git a/src/paimon/common/utils/arrow/mem_utils.cpp b/src/paimon/common/utils/arrow/mem_utils.cpp index 1e9695332..f454bda25 100644 --- a/src/paimon/common/utils/arrow/mem_utils.cpp +++ b/src/paimon/common/utils/arrow/mem_utils.cpp @@ -37,7 +37,7 @@ namespace { struct ArrowArrayPrivateData { void (*release)(ArrowArray*); void* private_data; - std::shared_ptr arrow_pool; + std::shared_ptr lifetime; }; void ReleaseArrowArray(ArrowArray* array) { @@ -48,7 +48,9 @@ void ReleaseArrowArray(ArrowArray* array) { array->release(array); } -} // namespace +bool HasSameOwner(const std::shared_ptr& lhs, const std::shared_ptr& rhs) { + return !lhs.owner_before(rhs) && !rhs.owner_before(lhs); +} class ArrowMemPoolAdaptor : public arrow::MemoryPool { public: @@ -122,30 +124,40 @@ class ArrowMemPoolAdaptor : public arrow::MemoryPool { arrow::internal::MemoryPoolStats stats_; }; +} // namespace + std::unique_ptr GetArrowPool(const std::shared_ptr& pool) { return std::make_unique(pool); } -Status RetainArrowArrayMemoryPool(ArrowArray* array, - const std::shared_ptr& arrow_pool) { - if (!array || !array->release) { - return Status::Invalid("cannot retain Arrow array memory pool"); +std::shared_ptr GetSharedArrowPool(const std::shared_ptr& pool) { + return std::make_shared(pool); +} + +Status AddArrowArrayLifetime(ArrowArray* array, const std::shared_ptr& lifetime) { + if (array == nullptr || array->release == nullptr) { + return Status::Invalid("cannot add lifetime to a released ArrowArray"); } - if (!arrow_pool) { - ArrowArrayRelease(array); - return Status::Invalid("cannot retain Arrow array memory pool"); + if (lifetime.use_count() == 0) { + return Status::Invalid("cannot add an empty lifetime to an ArrowArray"); + } + if (array->release == ReleaseArrowArray) { + const ArrowArrayPrivateData* data = + static_cast(array->private_data); + if (HasSameOwner(data->lifetime, lifetime)) { + return Status::OK(); + } } - std::unique_ptr data; try { - data = std::make_unique( - ArrowArrayPrivateData{array->release, array->private_data, arrow_pool}); + std::unique_ptr data = std::make_unique( + ArrowArrayPrivateData{array->release, array->private_data, lifetime}); + array->release = ReleaseArrowArray; + array->private_data = data.release(); + return Status::OK(); } catch (const std::bad_alloc&) { - ArrowArrayRelease(array); - return Status::OutOfMemory("failed to retain Arrow array memory pool"); + array->release(array); + return Status::OutOfMemory("failed to add lifetime to an ArrowArray"); } - array->private_data = data.release(); - array->release = ReleaseArrowArray; - return Status::OK(); } } // namespace paimon diff --git a/src/paimon/common/utils/arrow/mem_utils.h b/src/paimon/common/utils/arrow/mem_utils.h index 214bb4509..5f81437aa 100644 --- a/src/paimon/common/utils/arrow/mem_utils.h +++ b/src/paimon/common/utils/arrow/mem_utils.h @@ -21,6 +21,7 @@ #include +#include "arrow/c/abi.h" #include "arrow/memory_pool.h" #include "paimon/memory/memory_pool.h" #include "paimon/status.h" @@ -33,7 +34,15 @@ namespace paimon { PAIMON_EXPORT std::unique_ptr GetArrowPool( const std::shared_ptr& pool); -Status RetainArrowArrayMemoryPool(ArrowArray* array, - const std::shared_ptr& arrow_pool); +PAIMON_EXPORT std::shared_ptr GetSharedArrowPool( + const std::shared_ptr& pool); + +/// Keep an additional resource alive until the ArrowArray is released. +/// +/// The existing release callback and private data are preserved as a release chain, so this helper +/// can be applied to arrays produced by Arrow libraries or format plugins without inspecting their +/// private data. +PAIMON_EXPORT Status AddArrowArrayLifetime(ArrowArray* array, + const std::shared_ptr& lifetime); } // namespace paimon diff --git a/src/paimon/common/utils/arrow/mem_utils_test.cpp b/src/paimon/common/utils/arrow/mem_utils_test.cpp index 2e9bda22d..67878b716 100644 --- a/src/paimon/common/utils/arrow/mem_utils_test.cpp +++ b/src/paimon/common/utils/arrow/mem_utils_test.cpp @@ -22,10 +22,16 @@ #include #include #include +#include +#include "arrow/api.h" +#include "arrow/c/abi.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" #include "arrow/status.h" #include "gtest/gtest.h" #include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" namespace paimon::test { namespace { @@ -65,6 +71,59 @@ class FailingMemoryPool : public MemoryPool { Mode mode_; }; +class TrackingMemoryPool : public MemoryPool { + public: + TrackingMemoryPool(bool* destroyed, int32_t* free_count) + : destroyed_(destroyed), free_count_(free_count), delegate_(GetDefaultPool()) {} + + ~TrackingMemoryPool() override { + *destroyed_ = true; + } + + void* Malloc(uint64_t size, uint64_t alignment = 0) override { + return delegate_->Malloc(size, alignment); + } + + void* Realloc(void* p, size_t old_size, size_t new_size, uint64_t alignment = 0) override { + return delegate_->Realloc(p, old_size, new_size, alignment); + } + + void Free(void* p, uint64_t size) override { + ++*free_count_; + delegate_->Free(p, size); + } + + void Free(void* p, uint64_t size, uint64_t alignment) override { + ++*free_count_; + delegate_->Free(p, size, alignment); + } + + uint64_t CurrentUsage() const override { + return delegate_->CurrentUsage(); + } + + uint64_t MaxMemoryUsage() const override { + return delegate_->MaxMemoryUsage(); + } + + private: + bool* destroyed_; + int32_t* free_count_; + std::shared_ptr delegate_; +}; + +struct TrackingArrowArrayPrivateData { + std::shared_ptr lifetime; + std::vector* release_order; +}; + +void ReleaseTrackingArrowArray(ArrowArray* array) { + std::unique_ptr data( + static_cast(array->private_data)); + data->release_order->push_back(0); + array->release = nullptr; +} + } // namespace TEST(MemUtilsTest, TestSimple) { @@ -145,4 +204,88 @@ TEST(MemUtilsTest, TestReallocateOutOfMemory) { ASSERT_TRUE(throw_pool->Reallocate(/*old_size=*/0, /*new_size=*/16, 64, &ptr).IsOutOfMemory()); } +TEST(MemUtilsTest, TestAddArrowArrayLifetimeComposesReleaseChain) { + std::vector release_order; + std::shared_ptr original_lifetime(new int32_t(1), [&release_order](void* ptr) { + delete static_cast(ptr); + release_order.push_back(1); + }); + std::shared_ptr inner_lifetime(new int32_t(2), [&release_order](void* ptr) { + delete static_cast(ptr); + release_order.push_back(2); + }); + std::shared_ptr outer_lifetime(new int32_t(3), [&release_order](void* ptr) { + delete static_cast(ptr); + release_order.push_back(3); + }); + ArrowArray array{}; + array.release = ReleaseTrackingArrowArray; + array.private_data = new TrackingArrowArrayPrivateData{original_lifetime, &release_order}; + + ASSERT_OK(AddArrowArrayLifetime(&array, inner_lifetime)); + ASSERT_OK(AddArrowArrayLifetime(&array, outer_lifetime)); + original_lifetime.reset(); + inner_lifetime.reset(); + outer_lifetime.reset(); + + ArrowArrayRelease(&array); + ASSERT_EQ(nullptr, array.release); + ASSERT_EQ(std::vector({0, 1, 2, 3}), release_order); +} + +TEST(MemUtilsTest, TestAddArrowArrayLifetimeKeepsPoolAliveUntilOriginalRelease) { + bool pool_destroyed = false; + int32_t free_count = 0; + std::shared_ptr paimon_pool = + std::make_shared(&pool_destroyed, &free_count); + std::weak_ptr weak_paimon_pool = paimon_pool; + std::shared_ptr arrow_pool = GetSharedArrowPool(paimon_pool); + ArrowArray c_array{}; + { + arrow::Int32Builder builder(arrow_pool.get()); + ASSERT_TRUE(builder.Append(42).ok()); + arrow::Result> array_result = builder.Finish(); + ASSERT_TRUE(array_result.ok()) << array_result.status().ToString(); + std::shared_ptr array = std::move(array_result).ValueOrDie(); + ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); + } + ASSERT_OK(AddArrowArrayLifetime(&c_array, arrow_pool)); + + arrow_pool.reset(); + paimon_pool.reset(); + ASSERT_FALSE(weak_paimon_pool.expired()); + ASSERT_FALSE(pool_destroyed); + + ArrowArrayRelease(&c_array); + ASSERT_EQ(nullptr, c_array.release); + ASSERT_GT(free_count, 0); + ASSERT_TRUE(weak_paimon_pool.expired()); + ASSERT_TRUE(pool_destroyed); +} + +TEST(MemUtilsTest, TestAddArrowArrayLifetimeDeduplicatesSameOwner) { + std::vector release_order; + std::shared_ptr lifetime = std::make_shared(1); + ArrowArray array{}; + array.release = ReleaseTrackingArrowArray; + array.private_data = new TrackingArrowArrayPrivateData{nullptr, &release_order}; + + ASSERT_OK(AddArrowArrayLifetime(&array, lifetime)); + const long use_count = lifetime.use_count(); + ASSERT_OK(AddArrowArrayLifetime(&array, lifetime)); + ASSERT_EQ(use_count, lifetime.use_count()); + + ArrowArrayRelease(&array); +} + +TEST(MemUtilsTest, TestAddArrowArrayLifetimeRejectsInvalidInput) { + ArrowArray array{}; + ASSERT_NOK(AddArrowArrayLifetime(nullptr, std::make_shared(1))); + ASSERT_NOK(AddArrowArrayLifetime(&array, std::make_shared(1))); + + array.release = ReleaseTrackingArrowArray; + ASSERT_NOK(AddArrowArrayLifetime(&array, nullptr)); + array.release = nullptr; +} + } // namespace paimon::test diff --git a/src/paimon/core/append/append_only_writer_test.cpp b/src/paimon/core/append/append_only_writer_test.cpp index 114016091..bf64420e1 100644 --- a/src/paimon/core/append/append_only_writer_test.cpp +++ b/src/paimon/core/append/append_only_writer_test.cpp @@ -280,7 +280,8 @@ class AppendOnlyWriterTest : public testing::Test { auto c_file_schema = reader->GetFileSchema().value(); ASSERT_OK(reader->SetReadSchema(c_file_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(reader))); ASSERT_TRUE(expected_array->Equals(result_array)) << "Expected:\n" << expected_array->ToString() << "\nActual:\n" @@ -816,7 +817,8 @@ TEST_F(AppendOnlyWriterTest, TestWriteWithOnlyBlobField) { ASSERT_TRUE(arrow::ExportSchema(*schema, &c_blob_schema).ok()); ASSERT_OK(blob_reader->SetReadSchema(&c_blob_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto actual_array, ReadResultCollector::CollectResult(blob_reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual_array, + ReadResultCollector::CollectResult(std::move(blob_reader))); auto expected_struct_array = arrow::StructArray::Make({blob_array}, {blob_field->name()}).ValueOrDie(); auto expected_array = std::make_shared(expected_struct_array); diff --git a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp index 6c0d23be0..8cc1d800c 100644 --- a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp +++ b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp @@ -26,6 +26,7 @@ #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/read_ahead_cache.h" #include "paimon/executor.h" #include "paimon/testing/mock/mock_file_batch_reader.h" @@ -50,6 +51,7 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, target_type_ = arrow::struct_({arrow::field("f1", int_type_)}); pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); fs_ = std::make_shared(); ASSERT_OK_AND_ASSIGN(executor_, CreateDefaultExecutor(/*thread_count=*/2)); } @@ -88,7 +90,7 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_)); + /*enable_io_metrics=*/false, pool_, arrow_pool_)); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); @@ -115,6 +117,7 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, std::shared_ptr int_type_; std::shared_ptr target_type_; std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::shared_ptr fs_; std::shared_ptr executor_; }; diff --git a/src/paimon/core/io/async_key_value_projection_reader.h b/src/paimon/core/io/async_key_value_projection_reader.h index 79c7aa4de..7e9a8b3bb 100644 --- a/src/paimon/core/io/async_key_value_projection_reader.h +++ b/src/paimon/core/io/async_key_value_projection_reader.h @@ -33,10 +33,11 @@ class AsyncKeyValueProjectionReader : public BatchReader { const std::shared_ptr& target_schema, const std::vector& target_to_src_mapping, int32_t batch_size, int32_t projection_thread_num, - const std::shared_ptr& pool) { - auto create_consumer = [target_schema, target_to_src_mapping, pool]() + const std::shared_ptr& arrow_pool) { + auto create_consumer = [target_schema, target_to_src_mapping, arrow_pool]() -> Result>> { - return KeyValueProjectionConsumer::Create(target_schema, target_to_src_mapping, pool); + return KeyValueProjectionConsumer::Create(target_schema, target_to_src_mapping, + arrow_pool); }; std::unique_ptr producer = std::make_unique(std::move(sort_merge_reader), diff --git a/src/paimon/core/io/complete_row_tracking_fields_reader.cpp b/src/paimon/core/io/complete_row_tracking_fields_reader.cpp index f43ffcc07..3bd347677 100644 --- a/src/paimon/core/io/complete_row_tracking_fields_reader.cpp +++ b/src/paimon/core/io/complete_row_tracking_fields_reader.cpp @@ -35,11 +35,11 @@ namespace paimon { CompleteRowTrackingFieldsBatchReader::CompleteRowTrackingFieldsBatchReader( std::unique_ptr&& reader, const std::optional& first_row_id, int64_t snapshot_id, const std::optional>& file_field_names, - const std::shared_ptr& pool) + const std::shared_ptr& arrow_pool) : first_row_id_(first_row_id), snapshot_id_(snapshot_id), file_field_names_(file_field_names), - arrow_pool_(GetArrowPool(pool)), + arrow_pool_(arrow_pool), reader_(std::move(reader)) {} Status CompleteRowTrackingFieldsBatchReader::SetReadSchema( @@ -149,6 +149,7 @@ CompleteRowTrackingFieldsBatchReader::NextBatchWithBitmap() { arrow::StructArray::Make(sub_array_vec, read_schema_->field_names())); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*target_array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return std::move(src_array_with_bitmap); } diff --git a/src/paimon/core/io/complete_row_tracking_fields_reader.h b/src/paimon/core/io/complete_row_tracking_fields_reader.h index 7f4d7ead8..21ec2211c 100644 --- a/src/paimon/core/io/complete_row_tracking_fields_reader.h +++ b/src/paimon/core/io/complete_row_tracking_fields_reader.h @@ -39,7 +39,7 @@ class CompleteRowTrackingFieldsBatchReader : public FileBatchReader { CompleteRowTrackingFieldsBatchReader( std::unique_ptr&& reader, const std::optional& first_row_id, int64_t snapshot_id, const std::optional>& file_field_names, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result> GetFileSchema() const override { return Status::Invalid( diff --git a/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp b/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp index b6800c637..cad0a6c98 100644 --- a/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp +++ b/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp @@ -29,6 +29,7 @@ #include "arrow/array/array_nested.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/mock/mock_file_batch_reader.h" @@ -39,6 +40,7 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } void CheckResult(const std::shared_ptr& src_array, int64_t first_row_id, @@ -48,9 +50,9 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), batch_size); auto complete_row_tracking_fields_reader = - std::make_shared( + std::make_unique( std::move(file_batch_reader), first_row_id, snapshot_id, - /*file_field_names=*/std::nullopt, pool_); + /*file_field_names=*/std::nullopt, arrow_pool_); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK(complete_row_tracking_fields_reader->SetReadSchema( @@ -59,8 +61,7 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto result_with_special_fields, paimon::test::ReadResultCollector::CollectResult( - complete_row_tracking_fields_reader.get())); - complete_row_tracking_fields_reader->Close(); + std::move(complete_row_tracking_fields_reader))); auto expected_chunk_array = std::make_shared(expected_array); ASSERT_TRUE(result_with_special_fields->Equals(expected_chunk_array)); } @@ -76,7 +77,7 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/10, /*snapshot_id=*/1, - file_field_names, pool_); + file_field_names, arrow_pool_); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( @@ -92,6 +93,7 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { private: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestSetReadSchema) { @@ -378,7 +380,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidWithReadNonExistFiel auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/100, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, pool_); + /*file_field_names=*/std::nullopt, arrow_pool_); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( @@ -415,7 +417,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidNextBatchBeforeSetRe auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/100, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, pool_); + /*file_field_names=*/std::nullopt, arrow_pool_); ASSERT_NOK_WITH_MSG(complete_row_tracking_fields_reader->NextBatchWithBitmap(), "in CompleteRowTrackingFieldsBatchReader SetReadSchema is supposed to be " "called before NextBatch"); @@ -446,7 +448,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidNullFirstRowId) { auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/std::nullopt, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, pool_); + /*file_field_names=*/std::nullopt, arrow_pool_); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( diff --git a/src/paimon/core/io/field_mapping_reader.cpp b/src/paimon/core/io/field_mapping_reader.cpp index 2d07cef24..3a0986393 100644 --- a/src/paimon/core/io/field_mapping_reader.cpp +++ b/src/paimon/core/io/field_mapping_reader.cpp @@ -143,9 +143,9 @@ Result> FieldMappingReader::Create( int32_t field_count, std::unique_ptr&& reader, const BinaryRow& partition, std::unique_ptr&& mapping, std::set&& skip_map_selected_keys_filter_field_ids, - const std::shared_ptr& pool) { + const std::shared_ptr& arrow_pool) { auto mapping_reader = std::unique_ptr(new FieldMappingReader( - field_count, std::move(reader), partition, std::move(mapping), pool)); + field_count, std::move(reader), partition, std::move(mapping), arrow_pool)); mapping_reader->need_mapping_ = false; mapping_reader->need_casting_ = false; @@ -202,9 +202,9 @@ FieldMappingReader::FieldMappingReader(int32_t field_count, std::unique_ptr&& reader, const BinaryRow& partition, std::unique_ptr&& mapping, - const std::shared_ptr& pool) + const std::shared_ptr& arrow_pool) : field_count_(field_count), - arrow_pool_(GetArrowPool(pool)), + arrow_pool_(arrow_pool), reader_(std::move(reader)), partition_(partition), partition_info_(mapping->partition_info), @@ -321,6 +321,7 @@ Result FieldMappingReader::NextBatchWithBitmap std::unique_ptr target_c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*arrow_array, target_c_arrow_array.get(), target_c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(target_c_arrow_array.get(), arrow_pool_)); auto target_batch = std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema)); return std::make_pair(std::move(target_batch), std::move(bitmap)); } diff --git a/src/paimon/core/io/field_mapping_reader.h b/src/paimon/core/io/field_mapping_reader.h index 7a5edc16a..dd3d8526e 100644 --- a/src/paimon/core/io/field_mapping_reader.h +++ b/src/paimon/core/io/field_mapping_reader.h @@ -53,7 +53,7 @@ class FieldMappingReader : public FileBatchReader { int32_t field_count, std::unique_ptr&& reader, const BinaryRow& partition, std::unique_ptr&& mapping, std::set&& skip_map_selected_keys_filter_field_ids, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result NextBatch() override { return Status::Invalid( @@ -94,7 +94,7 @@ class FieldMappingReader : public FileBatchReader { private: FieldMappingReader(int32_t field_count, std::unique_ptr&& reader, const BinaryRow& partition, std::unique_ptr&& mapping, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result> GenerateSinglePartitionArray(int32_t idx, int32_t batch_size) const; diff --git a/src/paimon/core/io/field_mapping_reader_test.cpp b/src/paimon/core/io/field_mapping_reader_test.cpp index 5375a250f..9a5a80bfa 100644 --- a/src/paimon/core/io/field_mapping_reader_test.cpp +++ b/src/paimon/core/io/field_mapping_reader_test.cpp @@ -35,6 +35,7 @@ #include "gtest/gtest.h" #include "paimon/common/data/blob_utils.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/utils/field_mapping.h" #include "paimon/defs.h" @@ -60,6 +61,7 @@ namespace paimon::test { class FieldMappingReaderTest : public ::testing::Test { public: void SetUp() override { + arrow_pool_ = GetSharedArrowPool(pool_); partition_ = BinaryRowGenerator::GenerateRow({10, 0}, pool_.get()); } void TearDown() override {} @@ -139,8 +141,9 @@ class FieldMappingReaderTest : public ::testing::Test { auto reader, FieldMappingReader::Create( /*field_count=*/read_schema->num_fields(), std::move(orc_batch_reader), partition_, - std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, pool_)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(reader))); if (expect_array == nullptr && result_array == nullptr) { // expect empty result return; @@ -180,8 +183,9 @@ class FieldMappingReaderTest : public ::testing::Test { auto reader, FieldMappingReader::Create( /*field_count=*/read_schema->num_fields(), std::move(orc_batch_reader), partition, - std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, pool_)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(reader))); if (expect_array == nullptr && result_array == nullptr) { // expect empty result return; @@ -198,6 +202,7 @@ class FieldMappingReaderTest : public ::testing::Test { private: std::shared_ptr pool_ = GetDefaultPool(); + std::shared_ptr arrow_pool_; std::vector data_fields_ = {DataField(0, arrow::field("f0", arrow::utf8())), DataField(1, arrow::field("f1", arrow::int32())), DataField(2, arrow::field("f2", arrow::int32())), @@ -242,7 +247,7 @@ TEST_F(FieldMappingReaderTest, TestGenerateSinglePartitionArray) { auto mapping_reader, FieldMappingReader::Create( /*field_count=*/8, /*reader=*/nullptr, partition, std::move(field_mapping), - /*skip_map_selected_keys_filter_field_ids=*/{}, pool_)); + /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_)); { ASSERT_OK_AND_ASSIGN(auto p7_array, mapping_reader->GenerateSinglePartitionArray( @@ -470,11 +475,12 @@ TEST_F(FieldMappingReaderTest, TestSchemaEvolutionAddedFieldInsideList) { ASSERT_OK_AND_ASSIGN(auto mapping, mapping_builder->CreateFieldMapping(data_fields)); auto mock = std::make_unique( data_array, arrow::struct_(data_schema->fields()), /*read_batch_size=*/8); - ASSERT_OK_AND_ASSIGN(auto reader, FieldMappingReader::Create( - read_schema->num_fields(), std::move(mock), - BinaryRow::EmptyRow(), std::move(mapping), - /*skip_map_selected_keys_filter_field_ids=*/{}, pool_)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(reader.get())); + ASSERT_OK_AND_ASSIGN( + auto reader, + FieldMappingReader::Create(read_schema->num_fields(), std::move(mock), + BinaryRow::EmptyRow(), std::move(mapping), + /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto expect_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(read_schema->fields()), R"([ @@ -863,7 +869,7 @@ TEST_F(FieldMappingReaderTest, TestCreateFailFastOnInvalidMapSelectedKeysMetadat /*field_count=*/1, /*reader=*/nullptr, /*partition=*/BinaryRow::EmptyRow(), std::move(field_mapping), - /*skip_map_selected_keys_filter_field_ids=*/{}, pool_), + /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_), "Duplicate selected key 'a'"); } } // namespace paimon::test diff --git a/src/paimon/core/io/generic_row_to_arrow_array_converter.cpp b/src/paimon/core/io/generic_row_to_arrow_array_converter.cpp index ee442e64e..385e59986 100644 --- a/src/paimon/core/io/generic_row_to_arrow_array_converter.cpp +++ b/src/paimon/core/io/generic_row_to_arrow_array_converter.cpp @@ -32,10 +32,10 @@ namespace paimon { Result> GenericRowToArrowArrayConverter::Create( - const std::shared_ptr& schema, arrow::MemoryPool* pool) { + const std::shared_ptr& schema, const std::shared_ptr& pool) { std::unique_ptr array_builder; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder( - pool, std::make_shared(schema->fields()), &array_builder)); + pool.get(), std::make_shared(schema->fields()), &array_builder)); auto struct_builder = checked_pointer_cast(std::move(array_builder)); std::vector appenders; @@ -48,7 +48,7 @@ Result> GenericRowToArrowArrayC appenders.emplace_back(std::move(func)); } return std::unique_ptr(new GenericRowToArrowArrayConverter( - reserve_count, std::move(appenders), std::move(struct_builder), nullptr)); + reserve_count, std::move(appenders), std::move(struct_builder), pool)); } Result GenericRowToArrowArrayConverter::NextBatch( diff --git a/src/paimon/core/io/generic_row_to_arrow_array_converter.h b/src/paimon/core/io/generic_row_to_arrow_array_converter.h index 77df6c8f9..9517de37c 100644 --- a/src/paimon/core/io/generic_row_to_arrow_array_converter.h +++ b/src/paimon/core/io/generic_row_to_arrow_array_converter.h @@ -40,16 +40,17 @@ class GenericRowToArrowArrayConverter : public RowToArrowArrayConverter { public: static Result> Create( - const std::shared_ptr& schema, arrow::MemoryPool* pool); + const std::shared_ptr& schema, + const std::shared_ptr& pool); Result NextBatch(const std::vector& rows) override; private: GenericRowToArrowArrayConverter(int32_t reserve_count, std::vector&& appenders, std::unique_ptr&& array_builder, - std::unique_ptr&& arrow_pool) + const std::shared_ptr& arrow_pool) : RowToArrowArrayConverter(reserve_count, std::move(appenders), std::move(array_builder), - std::move(arrow_pool)) {} + arrow_pool) {} }; } // namespace paimon diff --git a/src/paimon/core/io/key_value_meta_projection_consumer.cpp b/src/paimon/core/io/key_value_meta_projection_consumer.cpp index e7e5b854d..7e9117412 100644 --- a/src/paimon/core/io/key_value_meta_projection_consumer.cpp +++ b/src/paimon/core/io/key_value_meta_projection_consumer.cpp @@ -60,7 +60,7 @@ Result> KeyValueMetaProjectionCo target_to_src_mapping.size())); } - auto arrow_pool = GetArrowPool(pool); + std::shared_ptr arrow_pool = GetSharedArrowPool(pool); // target fields of output array: special fields + value fields std::unique_ptr array_builder; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder( @@ -93,7 +93,7 @@ Result> KeyValueMetaProjectionCo appenders.emplace_back(func); } return std::unique_ptr(new KeyValueMetaProjectionConsumer( - reserve_count, std::move(appenders), std::move(struct_builder), std::move(arrow_pool), + reserve_count, std::move(appenders), std::move(struct_builder), arrow_pool, target_to_src_mapping, sequence_appender, value_kind_appender)); } diff --git a/src/paimon/core/io/key_value_meta_projection_consumer.h b/src/paimon/core/io/key_value_meta_projection_consumer.h index ac67aa49b..5c42e7aed 100644 --- a/src/paimon/core/io/key_value_meta_projection_consumer.h +++ b/src/paimon/core/io/key_value_meta_projection_consumer.h @@ -55,12 +55,12 @@ class KeyValueMetaProjectionConsumer : public RowToArrowArrayConverter&& appenders, std::unique_ptr&& array_builder, - std::unique_ptr&& arrow_pool, + const std::shared_ptr& arrow_pool, const std::vector& target_to_src_mapping, arrow::Int64Builder* sequence_appender, arrow::Int8Builder* value_kind_appender) : RowToArrowArrayConverter(reserve_count, std::move(appenders), std::move(array_builder), - std::move(arrow_pool)), + arrow_pool), target_to_src_mapping_(target_to_src_mapping), sequence_appender_(sequence_appender), value_kind_appender_(value_kind_appender) {} diff --git a/src/paimon/core/io/key_value_projection_consumer.cpp b/src/paimon/core/io/key_value_projection_consumer.cpp index 53cbe0db7..c1191de13 100644 --- a/src/paimon/core/io/key_value_projection_consumer.cpp +++ b/src/paimon/core/io/key_value_projection_consumer.cpp @@ -38,12 +38,12 @@ class MemoryPool; Result> KeyValueProjectionConsumer::Create( const std::shared_ptr& target_schema, - const std::vector& target_to_src_mapping, const std::shared_ptr& pool) { + const std::vector& target_to_src_mapping, + const std::shared_ptr& arrow_pool) { if (static_cast(target_schema->num_fields()) != target_to_src_mapping.size()) { return Status::Invalid( "target_schema and target_to_src_mapping mismatch in KeyValueProjectionConsumer"); } - auto arrow_pool = GetArrowPool(pool); std::unique_ptr array_builder; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder( arrow_pool.get(), std::make_shared(target_schema->fields()), @@ -61,7 +61,7 @@ Result> KeyValueProjectionConsumer:: appenders.emplace_back(func); } return std::unique_ptr(new KeyValueProjectionConsumer( - reserve_count, std::move(appenders), std::move(struct_builder), std::move(arrow_pool), + reserve_count, std::move(appenders), std::move(struct_builder), arrow_pool, target_to_src_mapping)); } diff --git a/src/paimon/core/io/key_value_projection_consumer.h b/src/paimon/core/io/key_value_projection_consumer.h index 07020335e..3f16868ea 100644 --- a/src/paimon/core/io/key_value_projection_consumer.h +++ b/src/paimon/core/io/key_value_projection_consumer.h @@ -48,7 +48,8 @@ class KeyValueProjectionConsumer static Result> Create( const std::shared_ptr& target_schema, - const std::vector& target_to_src_mapping, const std::shared_ptr& pool); + const std::vector& target_to_src_mapping, + const std::shared_ptr& arrow_pool); Result NextBatch(const std::vector& key_value_vec) override; @@ -57,10 +58,10 @@ class KeyValueProjectionConsumer private: KeyValueProjectionConsumer(int32_t reserve_count, std::vector&& appenders, std::unique_ptr&& array_builder, - std::unique_ptr&& arrow_pool, + const std::shared_ptr& arrow_pool, const std::vector& target_to_src_mapping) : RowToArrowArrayConverter(reserve_count, std::move(appenders), std::move(array_builder), - std::move(arrow_pool)), + arrow_pool), target_to_src_mapping_(target_to_src_mapping) {} std::vector target_to_src_mapping_; diff --git a/src/paimon/core/io/key_value_projection_reader.cpp b/src/paimon/core/io/key_value_projection_reader.cpp index 72ada014e..a1f9ac2c6 100644 --- a/src/paimon/core/io/key_value_projection_reader.cpp +++ b/src/paimon/core/io/key_value_projection_reader.cpp @@ -36,10 +36,11 @@ Result> KeyValueProjectionReader::Crea std::unique_ptr&& sort_merge_reader, const std::shared_ptr& target_schema, const std::vector& target_to_src_mapping, int32_t batch_size, - const std::shared_ptr& pool) { + const std::shared_ptr& arrow_pool) { std::unique_ptr> projection_consumer; - PAIMON_ASSIGN_OR_RAISE(projection_consumer, KeyValueProjectionConsumer::Create( - target_schema, target_to_src_mapping, pool)); + PAIMON_ASSIGN_OR_RAISE( + projection_consumer, + KeyValueProjectionConsumer::Create(target_schema, target_to_src_mapping, arrow_pool)); return std::unique_ptr(new KeyValueProjectionReader( batch_size, std::move(sort_merge_reader), std::move(projection_consumer))); } diff --git a/src/paimon/core/io/key_value_projection_reader.h b/src/paimon/core/io/key_value_projection_reader.h index fc0fbc48f..301f7b923 100644 --- a/src/paimon/core/io/key_value_projection_reader.h +++ b/src/paimon/core/io/key_value_projection_reader.h @@ -44,7 +44,7 @@ class KeyValueProjectionReader : public BatchReader { std::unique_ptr&& sort_merge_reader, const std::shared_ptr& target_schema, const std::vector& target_to_src_mapping, int32_t batch_size, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); Result NextBatch() override; diff --git a/src/paimon/core/io/key_value_projection_reader_test.cpp b/src/paimon/core/io/key_value_projection_reader_test.cpp index a37711edb..6121698c1 100644 --- a/src/paimon/core/io/key_value_projection_reader_test.cpp +++ b/src/paimon/core/io/key_value_projection_reader_test.cpp @@ -58,6 +58,7 @@ class KeyValueProjectionReaderTest : public testing::Test, public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } std::unique_ptr GenerateProjectionReader( @@ -94,15 +95,16 @@ class KeyValueProjectionReaderTest : public testing::Test, std::move(concat_readers), user_key_comparator, /*user_defined_seq_comparator=*/nullptr, merge_function_wrapper); if (!multi_thread_row_to_batch) { - EXPECT_OK_AND_ASSIGN(auto projection_reader, KeyValueProjectionReader::Create( - std::move(sort_merge_reader), - target_schema, target_to_src_mapping, - /*batch_size=*/batch_size, pool_)); + EXPECT_OK_AND_ASSIGN( + auto projection_reader, + KeyValueProjectionReader::Create(std::move(sort_merge_reader), target_schema, + target_to_src_mapping, + /*batch_size=*/batch_size, arrow_pool_)); return std::move(projection_reader); } else { return std::make_unique( std::move(sort_merge_reader), target_schema, target_to_src_mapping, batch_size, - /*projection_thread_num=*/3, pool_); + /*projection_thread_num=*/3, arrow_pool_); } } @@ -157,6 +159,7 @@ class KeyValueProjectionReaderTest : public testing::Test, private: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; TEST_P(KeyValueProjectionReaderTest, TestBulkData) { @@ -545,8 +548,9 @@ TEST_P(KeyValueProjectionReaderTest, TestInvalidProducer) { auto projection_reader = GenerateProjectionReader(src_array, target_schema, target_to_src_mapping, key_schema, value_schema, /*batch_size=*/1, multi_thread_row_to_batch); - ASSERT_NOK_WITH_MSG(paimon::test::ReadResultCollector::CollectResult(projection_reader.get()), - "cannot cast VALUE_KIND column to int8 arrow array"); + ASSERT_NOK_WITH_MSG( + paimon::test::ReadResultCollector::CollectResult(std::move(projection_reader)), + "cannot cast VALUE_KIND column to int8 arrow array"); } INSTANTIATE_TEST_SUITE_P(EnableMultiThreadRowToBatch, KeyValueProjectionReaderTest, diff --git a/src/paimon/core/io/meta_to_arrow_array_converter.cpp b/src/paimon/core/io/meta_to_arrow_array_converter.cpp index c149a0e6e..1980a3478 100644 --- a/src/paimon/core/io/meta_to_arrow_array_converter.cpp +++ b/src/paimon/core/io/meta_to_arrow_array_converter.cpp @@ -27,7 +27,7 @@ Result> MetaToArrowArrayConverter::Cr return Status::Invalid("meta_data_type in MetaToArrowArrayConverter must be struct type"); } auto struct_type = checked_pointer_cast(meta_data_type); - auto arrow_pool = GetArrowPool(pool); + std::shared_ptr arrow_pool = GetSharedArrowPool(pool); std::unique_ptr array_builder; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder( arrow_pool.get(), arrow::struct_(struct_type->fields()), &array_builder)); @@ -44,7 +44,7 @@ Result> MetaToArrowArrayConverter::Cr appenders.emplace_back(func); } return std::unique_ptr(new MetaToArrowArrayConverter( - reserve_count, std::move(appenders), std::move(struct_builder), std::move(arrow_pool))); + reserve_count, std::move(appenders), std::move(struct_builder), arrow_pool)); } Result> MetaToArrowArrayConverter::NextBatch( diff --git a/src/paimon/core/io/meta_to_arrow_array_converter.h b/src/paimon/core/io/meta_to_arrow_array_converter.h index ad774e032..3ded061a0 100644 --- a/src/paimon/core/io/meta_to_arrow_array_converter.h +++ b/src/paimon/core/io/meta_to_arrow_array_converter.h @@ -50,8 +50,8 @@ class MetaToArrowArrayConverter private: MetaToArrowArrayConverter(int32_t reserve_count, std::vector&& appenders, std::unique_ptr&& array_builder, - std::unique_ptr&& arrow_pool) + const std::shared_ptr& arrow_pool) : RowToArrowArrayConverter(reserve_count, std::move(appenders), std::move(array_builder), - std::move(arrow_pool)) {} + arrow_pool) {} }; } // namespace paimon diff --git a/src/paimon/core/io/row_to_arrow_array_converter.h b/src/paimon/core/io/row_to_arrow_array_converter.h index 245336856..b4153d968 100644 --- a/src/paimon/core/io/row_to_arrow_array_converter.h +++ b/src/paimon/core/io/row_to_arrow_array_converter.h @@ -52,7 +52,7 @@ class RowToArrowArrayConverter { std::function; RowToArrowArrayConverter(int32_t reserve_count, std::vector&& appenders, std::unique_ptr&& array_builder, - std::unique_ptr&& arrow_pool); + const std::shared_ptr& arrow_pool); static Result AppendField(bool use_view, arrow::ArrayBuilder* array_builder, int32_t* reserve_count); @@ -72,7 +72,7 @@ class RowToArrowArrayConverter { protected: std::vector reserved_sizes_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::vector appenders_; std::unique_ptr array_builder_; }; @@ -86,9 +86,9 @@ template RowToArrowArrayConverter::RowToArrowArrayConverter( int32_t reserve_count, std::vector::AppendValueFunc>&& appenders, std::unique_ptr&& array_builder, - std::unique_ptr&& arrow_pool) + const std::shared_ptr& arrow_pool) : reserved_sizes_(reserve_count, -1), - arrow_pool_(std::move(arrow_pool)), + arrow_pool_(arrow_pool), appenders_(std::move(appenders)), array_builder_(std::move(array_builder)) {} @@ -110,6 +110,7 @@ Result RowToArrowArrayConverter::FinishAndAccumula std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return make_pair(std::move(c_array), std::move(c_schema)); } diff --git a/src/paimon/core/io/vector_file_batch_reader.cpp b/src/paimon/core/io/vector_file_batch_reader.cpp index f16025c42..518958307 100644 --- a/src/paimon/core/io/vector_file_batch_reader.cpp +++ b/src/paimon/core/io/vector_file_batch_reader.cpp @@ -220,8 +220,8 @@ Result> ConvertToReadType( } // namespace VectorFileBatchReader::VectorFileBatchReader(std::unique_ptr&& reader, - const std::shared_ptr& pool) - : arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)) {} + const std::shared_ptr& arrow_pool) + : arrow_pool_(arrow_pool), reader_(std::move(reader)) {} bool VectorFileBatchReader::ContainsVector(const std::shared_ptr& schema) { return VectorUtils::ContainsVector(schema); @@ -264,6 +264,7 @@ Result VectorFileBatchReader::ConvertBatch(ReadBatch&& b PAIMON_ASSIGN_OR_RAISE(array, ConvertToReadType(array, read_type_, arrow_pool_.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW(array->Validate()); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return std::move(batch); } diff --git a/src/paimon/core/io/vector_file_batch_reader.h b/src/paimon/core/io/vector_file_batch_reader.h index b7eab263c..29b064d87 100644 --- a/src/paimon/core/io/vector_file_batch_reader.h +++ b/src/paimon/core/io/vector_file_batch_reader.h @@ -39,7 +39,7 @@ class MemoryPool; class VectorFileBatchReader : public FileBatchReader { public: VectorFileBatchReader(std::unique_ptr&& reader, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); static bool ContainsVector(const std::shared_ptr& schema); diff --git a/src/paimon/core/io/vector_file_batch_reader_test.cpp b/src/paimon/core/io/vector_file_batch_reader_test.cpp index e334e62aa..5eff19e73 100644 --- a/src/paimon/core/io/vector_file_batch_reader_test.cpp +++ b/src/paimon/core/io/vector_file_batch_reader_test.cpp @@ -27,6 +27,7 @@ #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/mock/mock_file_batch_reader.h" @@ -61,7 +62,7 @@ TEST(VectorFileBatchReaderTest, ConvertSchemaAndNextBatch) { std::make_unique(physical_array, physical_type, /*batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); MockFileBatchReader* inner_reader = mock_reader.get(); - VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); ASSERT_TRUE(VectorFileBatchReader::ContainsVector(arrow::schema(logical_type->fields()))); ASSERT_FALSE(VectorFileBatchReader::ContainsVector(arrow::schema(physical_type->fields()))); @@ -97,7 +98,7 @@ TEST(VectorFileBatchReaderTest, KeepFixedSizeListFileSchema) { std::make_unique(logical_array, logical_type, /*batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); MockFileBatchReader* inner_reader = mock_reader.get(); - VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); @@ -135,7 +136,7 @@ TEST(VectorFileBatchReaderTest, NormalizeFixedSizeListElementField) { auto mock_reader = std::make_unique(file_array, file_type, /*batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); - VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); @@ -173,7 +174,7 @@ TEST(VectorFileBatchReaderTest, ConvertNestedVectorsWithBitmap) { auto mock_reader = std::make_unique(physical_array, physical_type, bitmap, /*read_batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); - VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, @@ -202,7 +203,7 @@ TEST(VectorFileBatchReaderTest, RejectInvalidVectorValues) { auto mock_reader = std::make_unique(physical_array, physical_type, /*read_batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); - VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE( arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); @@ -221,7 +222,7 @@ TEST(VectorFileBatchReaderTest, RejectInvalidFixedSizeListVectorValues) { auto mock_reader = std::make_unique(physical_array, physical_type, /*read_batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); - VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(physical_type->fields()), &c_read_schema).ok()); ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, diff --git a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp index 1a818dd01..5ca57f91b 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp @@ -21,6 +21,7 @@ #include #include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/mergetree/compact/first_row_merge_function_wrapper.h" #include "paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" @@ -84,7 +85,8 @@ LookupMergeTreeCompactRewriter::Create( read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr internal_context, - InternalReadContext::Create(read_context, table_schema, options.ToMap())); + InternalReadContext::Create(read_context, table_schema, options.ToMap(), + GetSharedArrowPool(read_context->GetMemoryPool()))); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr path_factory, path_factory_cache->GetOrCreatePathFactory(options.GetFileFormat()->Identifier())); diff --git a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp index c930c5ca7..bcaf72798 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp +++ b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp @@ -331,7 +331,7 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParamOpen(compact_file_name)); ASSERT_OK_AND_ASSIGN(auto file_batch_reader, reader_builder->Build(input_stream)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(file_batch_reader.get())); + ReadResultCollector::CollectResult(std::move(file_batch_reader))); // handle type nullable, as result_array does not have not null flag result_array = result_array->View(expected_array->type()).ValueOrDie(); diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp index 37c1cfa04..e81259cbd 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp @@ -24,6 +24,7 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/key_value_data_file_writer_factories.h" #include "paimon/core/io/key_value_meta_projection_consumer.h" @@ -84,7 +85,8 @@ Result> MergeTreeCompactRewriter::Crea read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr internal_context, - InternalReadContext::Create(read_context, table_schema, options.ToMap())); + InternalReadContext::Create(read_context, table_schema, options.ToMap(), + GetSharedArrowPool(read_context->GetMemoryPool()))); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr path_factory, path_factory_cache->GetOrCreatePathFactory(options.GetFileFormat()->Identifier())); diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter_test.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter_test.cpp index dc19c4726..27424a1ae 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter_test.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter_test.cpp @@ -88,7 +88,7 @@ class MergeTreeCompactRewriterTest : public testing::Test { fs->Open(compact_file_name)); ASSERT_OK_AND_ASSIGN(auto file_batch_reader, reader_builder->Build(input_stream)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(file_batch_reader.get())); + ReadResultCollector::CollectResult(std::move(file_batch_reader))); // handle type nullable, as result_array does not have not null flag result_array = result_array->View(expected_array->type()).ValueOrDie(); diff --git a/src/paimon/core/mergetree/lookup_levels.cpp b/src/paimon/core/mergetree/lookup_levels.cpp index 512f28728..7743d643c 100644 --- a/src/paimon/core/mergetree/lookup_levels.cpp +++ b/src/paimon/core/mergetree/lookup_levels.cpp @@ -20,6 +20,7 @@ #include "fmt/format.h" #include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/key_value_data_file_record_reader.h" #include "paimon/core/mergetree/lookup/file_position.h" @@ -68,7 +69,8 @@ Result>> LookupLevels::Create( read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr internal_read_context, - InternalReadContext::Create(read_context, table_schema, options.ToMap())); + InternalReadContext::Create(read_context, table_schema, options.ToMap(), + GetSharedArrowPool(read_context->GetMemoryPool()))); auto split_read = std::make_unique(path_factory, internal_read_context, pool, CreateDefaultExecutor()); diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index c1c804f9f..1885a976c 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -190,7 +190,7 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { file_format->CreateReaderBuilder(/*batch_size=*/10)); ASSERT_OK_AND_ASSIGN(auto orc_batch_reader, reader_builder->Build(input_stream)); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(orc_batch_reader.get())); + ReadResultCollector::CollectResult(std::move(orc_batch_reader))); ASSERT_TRUE(expected_array->Equals(result_array)) << result_array->ToString(); } diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index 5beee34ba..2f46f47ad 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -68,6 +68,7 @@ AbstractSplitRead::AbstractSplitRead(const std::shared_ptr const std::shared_ptr& memory_pool, const std::shared_ptr& executor) : pool_(memory_pool), + arrow_pool_(context->GetArrowMemoryPool()), executor_(executor), path_factory_(path_factory), options_(context->GetCoreOptions()), @@ -121,7 +122,7 @@ Result> AbstractSplitRead::ApplyPredicateFilterIfNe if (!context_->EnablePredicateFilter() || predicate == nullptr) { return std::move(reader); } - return PredicateBatchReader::Create(std::move(reader), predicate, pool_); + return PredicateBatchReader::Create(std::move(reader), predicate, arrow_pool_); } Result> AbstractSplitRead::PrepareReaderBuilder( @@ -154,8 +155,8 @@ Result> AbstractSplitRead::CreateFileBatchReade const std::string& file_format_identifier, const std::string& data_file_path, int64_t data_file_size, std::unique_ptr reader_builder) const { if (context_->EnableLateMaterializing()) { - reader_builder = - std::make_unique(std::move(reader_builder), pool_); + reader_builder = std::make_unique(std::move(reader_builder), + arrow_pool_); } // TODO(xinyu.lxy): test format table for mosaic format if (context_->EnablePrefetch() && file_format_identifier != "blob" && @@ -168,7 +169,8 @@ Result> AbstractSplitRead::CreateFileBatchReade context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), executor_, /*initialize_read_ranges=*/false, context_->ReadAheadCacheEnabled(), - context_->GetCacheConfig(), options_.PrefetchIoMetricsEnabled(), pool_)); + context_->GetCacheConfig(), options_.PrefetchIoMetricsEnabled(), pool_, + arrow_pool_)); return std::make_unique(std::move(prefetch_reader)); } else { PAIMON_ASSIGN_OR_RAISE( @@ -222,7 +224,7 @@ Result> AbstractSplitRead::CreateFieldMappingRe CreateFileBatchReader(file_format_identifier, data_file_path, file_meta->file_size, std::move(reader_builder))); if (VectorFileBatchReader::ContainsVector(read_schema)) { - file_reader = std::make_unique(std::move(file_reader), pool_); + file_reader = std::make_unique(std::move(file_reader), arrow_pool_); } std::set skip_map_selected_keys_filter_field_ids; if (file_format_identifier != "blob") { @@ -241,7 +243,7 @@ Result> AbstractSplitRead::CreateFieldMappingRe } file_reader = std::make_unique( std::move(file_reader), file_meta->first_row_id, file_meta->max_sequence_number, - file_field_names, pool_); + file_field_names, arrow_pool_); } const auto& predicate = field_mapping->non_partition_info.non_partition_filter; auto all_data_schema = DataField::ConvertDataFieldsToArrowSchema(data_schema->Fields()); @@ -254,11 +256,11 @@ Result> AbstractSplitRead::CreateFieldMappingRe return std::unique_ptr(); } - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr mapping_reader, - FieldMappingReader::Create(field_mapping_builder->GetReadFieldCount(), - std::move(final_reader), partition, std::move(field_mapping), - std::move(skip_map_selected_keys_filter_field_ids), pool_)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr mapping_reader, + FieldMappingReader::Create( + field_mapping_builder->GetReadFieldCount(), std::move(final_reader), + partition, std::move(field_mapping), + std::move(skip_map_selected_keys_filter_field_ids), arrow_pool_)); return mapping_reader; } @@ -326,8 +328,8 @@ AbstractSplitRead::ApplyShreddingReaderIfNeeded( } if (!plans.empty()) { - file_reader = - std::make_unique(std::move(file_reader), std::move(plans), pool_); + file_reader = std::make_unique(std::move(file_reader), + std::move(plans), arrow_pool_); } return std::make_pair(std::move(file_reader), std::move(handled_shared_shredding_field_ids)); } diff --git a/src/paimon/core/operation/abstract_split_read.h b/src/paimon/core/operation/abstract_split_read.h index a02ed5fb2..5a3f30683 100644 --- a/src/paimon/core/operation/abstract_split_read.h +++ b/src/paimon/core/operation/abstract_split_read.h @@ -42,6 +42,7 @@ #include "paimon/status.h" namespace arrow { +class MemoryPool; class Schema; } // namespace arrow @@ -129,6 +130,7 @@ class AbstractSplitRead : public SplitRead { protected: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::shared_ptr executor_; std::shared_ptr path_factory_; CoreOptions options_; diff --git a/src/paimon/core/operation/append_only_file_store_write.cpp b/src/paimon/core/operation/append_only_file_store_write.cpp index f660093a2..9a938f630 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -30,6 +30,7 @@ #include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/append/append_only_writer.h" #include "paimon/core/append/bucketed_append_compact_manager.h" @@ -290,8 +291,10 @@ Result> AppendOnlyFileStoreWrite::CreateFilesReader .WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, context_builder.Finish()); std::map options = options_.ToMap(); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_read_context, - InternalReadContext::Create(read_context, table_schema_, options)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr internal_read_context, + InternalReadContext::Create(read_context, table_schema_, options, + GetSharedArrowPool(read_context->GetMemoryPool()))); auto read = std::make_unique(file_store_path_factory_, internal_read_context, pool_, compact_executor_); diff --git a/src/paimon/core/operation/append_only_file_store_write_test.cpp b/src/paimon/core/operation/append_only_file_store_write_test.cpp index 4e1cef501..a6e4dc44a 100644 --- a/src/paimon/core/operation/append_only_file_store_write_test.cpp +++ b/src/paimon/core/operation/append_only_file_store_write_test.cpp @@ -188,7 +188,7 @@ class AppendOnlyFileStoreWriteTest : public testing::Test { EXPECT_OK(reader->SetReadSchema(c_file_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); EXPECT_OK_AND_ASSIGN(std::shared_ptr result, - ReadResultCollector::CollectResult(reader.get())); + ReadResultCollector::CollectResult(std::move(reader))); EXPECT_NE(nullptr, result); if (!result) { return nullptr; diff --git a/src/paimon/core/operation/data_evolution_split_read.cpp b/src/paimon/core/operation/data_evolution_split_read.cpp index 9e71fe47d..e6edbf0a5 100644 --- a/src/paimon/core/operation/data_evolution_split_read.cpp +++ b/src/paimon/core/operation/data_evolution_split_read.cpp @@ -219,7 +219,7 @@ Result> DataEvolutionSplitRead::CreateReader( InnerCreateReader(data_split, indexed_split->RowRanges())); if (HasIndexScoreField(raw_read_schema_)) { batch_reader = std::make_unique( - std::move(batch_reader), indexed_split->Scores(), pool_); + std::move(batch_reader), indexed_split->Scores(), arrow_pool_); } return WrapWithBlobViewResolverIfNeeded(data_split, std::move(batch_reader), indexed_split->RowRanges()); @@ -266,8 +266,9 @@ Result> DataEvolutionSplitRead::WrapWithBlobViewRes PAIMON_ASSIGN_OR_RAISE( BlobViewResolver resolver, BlobViewLookup::CreateResolver(blob_view_structs, catalog_context, pool_, executor)); - return std::make_unique( - std::move(inner_reader), std::move(read_blob_view_fields), std::move(resolver), pool_); + return std::make_unique(std::move(inner_reader), + std::move(read_blob_view_fields), + std::move(resolver), arrow_pool_); } Result> DataEvolutionSplitRead::CreateBlobViewReader( @@ -330,7 +331,7 @@ Result> DataEvolutionSplitRead::CreateBlobViewReade batch_readers.push_back(std::move(raw_file_reader)); } } - return std::make_unique(std::move(batch_readers), pool_); + return std::make_unique(std::move(batch_readers), arrow_pool_); } Result> DataEvolutionSplitRead::ExtractBlobViewStructs( @@ -438,11 +439,12 @@ Result> DataEvolutionSplitRead::InnerCreateReader( sub_readers.push_back(std::move(evolution_reader)); } } - auto concat_batch_reader = std::make_unique(std::move(sub_readers), pool_); + auto concat_batch_reader = + std::make_unique(std::move(sub_readers), arrow_pool_); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr batch_reader, ApplyPredicateFilterIfNeeded(std::move(concat_batch_reader), context_->GetPredicate())); - return std::make_unique(std::move(batch_reader), pool_); + return std::make_unique(std::move(batch_reader), arrow_pool_); } Result> DataEvolutionSplitRead::CreatePushDownPredicate( @@ -739,7 +741,7 @@ Result> DataEvolutionSplitRead::CreateU std::move(file_readers)); // Concat multiple blob files that map to the same data file. file_batch_readers[file_idx] = - std::make_unique(std::move(raw_readers), pool_); + std::make_unique(std::move(raw_readers), arrow_pool_); } } } @@ -747,7 +749,7 @@ Result> DataEvolutionSplitRead::CreateU // TODO(xinyu.lxy): check nullable when reader_offsets[read_field_idx] = -1 return DataEvolutionFileReader::Create(std::move(file_batch_readers), raw_read_schema_, options_.GetReadBatchSize(), reader_offsets, - field_offsets, pool_); + field_offsets, arrow_pool_); } namespace { @@ -849,9 +851,10 @@ Result> DataEvolutionSplitRead::CreateBlobFallbackR } groups.push_back(std::move(segments)); } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr fallback_reader, - BlobFallbackBatchReader::Create(std::move(groups), file_read_schema, - options_.GetReadBatchSize(), pool_)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr fallback_reader, + BlobFallbackBatchReader::Create(std::move(groups), file_read_schema, + options_.GetReadBatchSize(), arrow_pool_)); return std::move(fallback_reader); } diff --git a/src/paimon/core/operation/internal_read_context.cpp b/src/paimon/core/operation/internal_read_context.cpp index ecd70e0c8..e686c88ee 100644 --- a/src/paimon/core/operation/internal_read_context.cpp +++ b/src/paimon/core/operation/internal_read_context.cpp @@ -30,6 +30,7 @@ #include "paimon/common/predicate/predicate_validator.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/options/map_storage_layout.h" @@ -199,7 +200,8 @@ std::optional InternalReadContext::TryResolveSpecialFieldByName( Result> InternalReadContext::Create( const std::shared_ptr& context, const std::shared_ptr& table_schema, - const std::map& options) { + const std::map& options, + const std::shared_ptr& arrow_pool) { PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options, context->GetSpecificFileSystem(), context->GetFileSystemSchemeToIdentifierMap())); @@ -281,17 +283,22 @@ Result> InternalReadContext::Create( PredicateValidator::ValidatePredicateWithLiterals(context->GetPredicate())); } + if (!arrow_pool) { + return Status::Invalid("arrow memory pool is null"); + } return std::unique_ptr( - new InternalReadContext(context, table_schema, read_schema, core_options)); + new InternalReadContext(context, table_schema, read_schema, core_options, arrow_pool)); } InternalReadContext::InternalReadContext(const std::shared_ptr& read_context, const std::shared_ptr& table_schema, const std::shared_ptr& read_schema, - const CoreOptions& options) + const CoreOptions& options, + const std::shared_ptr& arrow_pool) : read_context_(read_context), table_schema_(table_schema), read_schema_(read_schema), + arrow_pool_(arrow_pool), options_(options) {} Result> InternalReadContext::CreateWithSchema( @@ -299,8 +306,9 @@ Result> InternalReadContext::CreateWithSche const std::shared_ptr& new_read_schema) { // Create a new InternalReadContext sharing all properties except read_schema. // The new read_schema is the minimal column set for COUNT(*). - return std::shared_ptr(new InternalReadContext( - original->read_context_, original->table_schema_, new_read_schema, original->options_)); + return std::shared_ptr( + new InternalReadContext(original->read_context_, original->table_schema_, new_read_schema, + original->options_, original->arrow_pool_)); } } // namespace paimon diff --git a/src/paimon/core/operation/internal_read_context.h b/src/paimon/core/operation/internal_read_context.h index 8e773cdcd..69c58236d 100644 --- a/src/paimon/core/operation/internal_read_context.h +++ b/src/paimon/core/operation/internal_read_context.h @@ -30,6 +30,7 @@ #include "paimon/result.h" namespace arrow { +class MemoryPool; class Schema; } // namespace arrow @@ -45,7 +46,8 @@ class InternalReadContext { static Result> Create( const std::shared_ptr& read_context, const std::shared_ptr& table_schema, - const std::map& options); + const std::map& options, + const std::shared_ptr& arrow_pool); const CoreOptions& GetCoreOptions() const { return options_; @@ -92,6 +94,9 @@ class InternalReadContext { std::shared_ptr GetMemoryPool() const { return read_context_->GetMemoryPool(); } + const std::shared_ptr& GetArrowMemoryPool() const { + return arrow_pool_; + } std::shared_ptr GetExecutor() const { return read_context_->GetExecutor(); } @@ -119,7 +124,8 @@ class InternalReadContext { InternalReadContext(const std::shared_ptr& read_context, const std::shared_ptr& table_schema, const std::shared_ptr& read_schema, - const CoreOptions& options); + const CoreOptions& options, + const std::shared_ptr& arrow_pool); static std::optional TryResolveSpecialFieldById(int32_t field_id, const CoreOptions& core_options); @@ -132,6 +138,7 @@ class InternalReadContext { std::shared_ptr read_context_; std::shared_ptr table_schema_; std::shared_ptr read_schema_; + std::shared_ptr arrow_pool_; CoreOptions options_; }; diff --git a/src/paimon/core/operation/internal_read_context_test.cpp b/src/paimon/core/operation/internal_read_context_test.cpp index 20838122b..7162da6fe 100644 --- a/src/paimon/core/operation/internal_read_context_test.cpp +++ b/src/paimon/core/operation/internal_read_context_test.cpp @@ -25,6 +25,7 @@ #include "gtest/gtest.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/data/shredding/map_shared_shredding_schema_utils.h" #include "paimon/defs.h" @@ -33,6 +34,18 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +Result> CreateInternalReadContextForTest( + const std::shared_ptr& read_context, + const std::shared_ptr& table_schema, + const std::map& options) { + return InternalReadContext::Create(read_context, table_schema, options, + GetSharedArrowPool(read_context->GetMemoryPool())); +} + +} // namespace + TEST(InternalReadContext, TestReadWithUnspecifiedSchema) { // no read schema is specified, read all fields std::string path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09"; @@ -41,8 +54,8 @@ TEST(InternalReadContext, TestReadWithUnspecifiedSchema) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + CreateInternalReadContextForTest(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(0, arrow::field("f0", arrow::utf8())), DataField(1, arrow::field("f1", arrow::int32())), DataField(2, arrow::field("f2", arrow::int32())), @@ -59,8 +72,8 @@ TEST(InternalReadContext, TestReadWithSpecifiedSchema) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + CreateInternalReadContextForTest(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8()))}; auto expected_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); @@ -75,8 +88,8 @@ TEST(InternalReadContext, TestReadWithSpecifiedFieldId) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + CreateInternalReadContextForTest(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8()))}; auto expected_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); @@ -94,8 +107,8 @@ TEST(InternalReadContext, TestReadWithSpecifiedFieldIdAndSchema) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + CreateInternalReadContextForTest(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8()))}; auto expected_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); @@ -117,7 +130,7 @@ TEST(InternalReadContext, TestReadWithRowTrackingAndScoreFields) { new_options[Options::DATA_EVOLUTION_ENABLED] = "true"; ASSERT_OK_AND_ASSIGN( auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, new_options)); + CreateInternalReadContextForTest(std::move(read_context), table_schema, new_options)); std::vector read_fields = { DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8())), SpecialFields::RowId(), @@ -133,8 +146,8 @@ TEST(InternalReadContext, TestReadWithRowTrackingAndScoreFields) { ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); - ASSERT_NOK_WITH_MSG(InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options()), + ASSERT_NOK_WITH_MSG(CreateInternalReadContextForTest(std::move(read_context), table_schema, + table_schema->Options()), "Get field _ROW_ID failed: not exist in table schema"); } { @@ -145,8 +158,8 @@ TEST(InternalReadContext, TestReadWithRowTrackingAndScoreFields) { ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); - ASSERT_NOK_WITH_MSG(InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options()), + ASSERT_NOK_WITH_MSG(CreateInternalReadContextForTest(std::move(read_context), table_schema, + table_schema->Options()), "Get field _INDEX_SCORE failed: not exist in table schema"); } } @@ -159,8 +172,8 @@ TEST(InternalReadContext, TestReadWithValueKindField) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + CreateInternalReadContextForTest(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), SpecialFields::ValueKind(), DataField(0, arrow::field("f0", arrow::utf8()))}; @@ -186,7 +199,7 @@ TEST(InternalReadContext, TestReadWithFieldIdsAndSpecialFields) { new_options[Options::DATA_EVOLUTION_ENABLED] = "true"; ASSERT_OK_AND_ASSIGN( auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, new_options)); + CreateInternalReadContextForTest(std::move(read_context), table_schema, new_options)); std::vector read_fields = { DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8())), SpecialFields::RowId(), @@ -215,7 +228,7 @@ TEST(InternalReadContext, TestReadWithProjectedSchemaAndSpecialFields) { ASSERT_OK_AND_ASSIGN(auto unique_read_context, context_builder.Finish()); std::shared_ptr read_context = std::move(unique_read_context); ASSERT_NOK_WITH_MSG( - InternalReadContext::Create(read_context, table_schema, table_schema->Options()), + CreateInternalReadContextForTest(read_context, table_schema, table_schema->Options()), "not exist in table schema"); } @@ -234,7 +247,7 @@ TEST(InternalReadContext, TestReadWithProjectedSchemaAndSpecialFields) { std::shared_ptr read_context = std::move(unique_read_context); ASSERT_OK_AND_ASSIGN( auto internal_context, - InternalReadContext::Create(read_context, table_schema, enabled_options)); + CreateInternalReadContextForTest(read_context, table_schema, enabled_options)); auto expected_schema = DataField::ConvertDataFieldsToArrowSchema(projected_fields); ASSERT_TRUE(internal_context->GetReadSchema()->Equals(expected_schema)); } @@ -258,7 +271,7 @@ TEST(InternalReadContext, TestReadWithProjectedSchemaWithoutFieldIds) { ASSERT_OK_AND_ASSIGN( auto internal_context, - InternalReadContext::Create(read_context, table_schema, table_schema->Options())); + CreateInternalReadContextForTest(read_context, table_schema, table_schema->Options())); std::vector expected_fields = { DataField(3, arrow::field("f3", arrow::float64())), @@ -290,7 +303,7 @@ TEST(InternalReadContext, TestProjectedSchemaMetadataWhitelist) { ASSERT_OK_AND_ASSIGN( auto internal_context, - InternalReadContext::Create(read_context, table_schema, table_schema->Options())); + CreateInternalReadContextForTest(read_context, table_schema, table_schema->Options())); auto aligned_field = internal_context->GetReadSchema()->GetFieldByName("f0"); ASSERT_TRUE(aligned_field); @@ -331,7 +344,7 @@ TEST(InternalReadContext, TestMapSharedShreddingAccessRequiresSharedShreddingLay std::shared_ptr read_context = std::move(unique_read_context); ASSERT_NOK_WITH_MSG( - InternalReadContext::Create(read_context, table_schema, table_schema->Options()), + CreateInternalReadContextForTest(read_context, table_schema, table_schema->Options()), "Selected-key MAP pushdown only supports top-level shared-shredding MAP field: tags"); } diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index 9b913edfb..4193fb37c 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -237,7 +237,7 @@ class MergeFileSplitRead::RealtimeReaderBuilder { if (record_readers.empty()) { record_readers_guard.Release(); return std::make_unique(std::vector>{}, - owner_->pool_); + owner_->arrow_pool_); } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, owner_->CreateSortMergeReader(std::move(record_readers))); @@ -335,7 +335,7 @@ Result> MergeFileSplitRead::CreateReader( } else { PAIMON_ASSIGN_OR_RAISE(batch_reader, CreateMergeReader(data_split, data_file_path_factory)); } - return std::make_unique(std::move(batch_reader), pool_); + return std::make_unique(std::move(batch_reader), arrow_pool_); } Result> MergeFileSplitRead::CreateRealtimeReader( @@ -435,7 +435,8 @@ Result> MergeFileSplitRead::CreateMergeReader( data_file_path_factory)); batch_readers.push_back(std::move(projection_reader)); } - auto concat_batch_reader = std::make_unique(std::move(batch_readers), pool_); + auto concat_batch_reader = + std::make_unique(std::move(batch_readers), arrow_pool_); return AbstractSplitRead::ApplyPredicateFilterIfNeeded(std::move(concat_batch_reader), context_->GetPredicate()); } @@ -463,7 +464,8 @@ Result> MergeFileSplitRead::CreateNoMergeReader( auto raw_readers = ObjectUtils::MoveVector>(std::move(raw_file_readers)); - auto concat_batch_reader = std::make_unique(std::move(raw_readers), pool_); + auto concat_batch_reader = + std::make_unique(std::move(raw_readers), arrow_pool_); return AbstractSplitRead::ApplyPredicateFilterIfNeeded(std::move(concat_batch_reader), context_->GetPredicate()); } @@ -684,13 +686,14 @@ Result> MergeFileSplitRead::CreateProjectedReader( PAIMON_ASSIGN_OR_RAISE( projection_reader, KeyValueProjectionReader::Create(std::move(sort_merge_reader), raw_read_schema_, - projection_, options_.GetReadBatchSize(), pool_)); + projection_, options_.GetReadBatchSize(), + arrow_pool_)); } else { const int32_t thread_number = context_->GetRowToBatchThreadNumber(); assert(thread_number > 0); projection_reader = std::make_unique( std::move(sort_merge_reader), raw_read_schema_, projection_, - options_.GetReadBatchSize(), thread_number, pool_); + options_.GetReadBatchSize(), thread_number, arrow_pool_); } ScopeGuard projection_reader_guard([&projection_reader]() { if (projection_reader) { @@ -702,7 +705,8 @@ Result> MergeFileSplitRead::CreateProjectedReader( projection_reader_guard.Release(); projection_reader = std::move(filtered_reader); if (complete_row_kind) { - return std::make_unique(std::move(projection_reader), pool_); + return std::make_unique(std::move(projection_reader), + arrow_pool_); } return projection_reader; } diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index d02120de4..4b7a7fab2 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -36,6 +36,7 @@ #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/utils/arrow/mem_utils.h" #include "paimon/common/utils/fields_comparator.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" @@ -93,8 +94,10 @@ class MergeFileSplitReadTest : public ::testing::Test, const std::shared_ptr& read_context, int32_t schema_id = 0) { SchemaManager schema_manager(fs_, read_context->GetPath()); EXPECT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(schema_id)); - EXPECT_OK_AND_ASSIGN(auto context, InternalReadContext::Create(read_context, table_schema, - read_context->GetOptions())); + EXPECT_OK_AND_ASSIGN( + auto context, + InternalReadContext::Create(read_context, table_schema, read_context->GetOptions(), + GetSharedArrowPool(read_context->GetMemoryPool()))); return context; } @@ -362,11 +365,12 @@ class MergeFileSplitReadTest : public ::testing::Test, split_read->CreateReader(split)); batch_readers.emplace_back(std::move(reader)); } - return std::make_unique(std::move(batch_readers), pool_); + return std::make_unique(std::move(batch_readers), arrow_pool_); } private: std::shared_ptr pool_ = GetDefaultPool(); + std::shared_ptr arrow_pool_ = GetSharedArrowPool(pool_); std::shared_ptr fs_ = std::make_shared(); std::shared_ptr executor_ = CreateDefaultExecutor(); }; @@ -649,7 +653,7 @@ TEST_P(MergeFileSplitReadTest, TestSimple) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, data_splits)); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -763,7 +767,7 @@ TEST_P(MergeFileSplitReadTest, TestLookUp) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -834,7 +838,7 @@ TEST_P(MergeFileSplitReadTest, TestDeduplicateMergeEngineWithDeleteMsg) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit2())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -893,7 +897,7 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicate) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -951,7 +955,7 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicateAndLateMaterializing) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1001,7 +1005,7 @@ TEST_P(MergeFileSplitReadTest, TestReadWithAlterTable) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1050,7 +1054,7 @@ TEST_P(MergeFileSplitReadTest, TestReadWithAlterTableWithReverseSequence) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1098,7 +1102,7 @@ TEST_P(MergeFileSplitReadTest, TestAggregateMergeEngine) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1145,7 +1149,7 @@ TEST_P(MergeFileSplitReadTest, TestPartialUpdateMergeEngine) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1190,7 +1194,7 @@ TEST_P(MergeFileSplitReadTest, TestPartialUpdateMergeEngineWithIgnoreDelete) { ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit2())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1230,7 +1234,7 @@ TEST_P(MergeFileSplitReadTest, TestPartialUpdateMergeEngineWithRemoveRecordOnDel ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit2())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_schema->fields(); fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1270,7 +1274,7 @@ TEST_P(MergeFileSplitReadTest, TestEmptyPlan) { std::vector> empty_data_split; ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, empty_data_split)); ASSERT_OK_AND_ASSIGN(std::shared_ptr read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); // empty result with null pointer batch ASSERT_FALSE(read_result); } @@ -1323,7 +1327,7 @@ TEST_P(MergeFileSplitReadTest, TestIOException) { io_hook->Reset(i, IOHook::Mode::RETURN_ERROR); auto batch_reader = CreateReader(internal_context, PrepareDataSplit()); CHECK_HOOK_STATUS(batch_reader.status(), i); - auto read_result = ReadResultCollector::CollectResult(batch_reader.value().get()); + auto read_result = ReadResultCollector::CollectResult(std::move(batch_reader).value()); CHECK_HOOK_STATUS(read_result.status(), i); auto result_array = read_result.value(); CheckResult(result_array, expected_array, read_schema); diff --git a/src/paimon/core/operation/raw_file_split_read.cpp b/src/paimon/core/operation/raw_file_split_read.cpp index 11fdec251..8439ba19f 100644 --- a/src/paimon/core/operation/raw_file_split_read.cpp +++ b/src/paimon/core/operation/raw_file_split_read.cpp @@ -137,10 +137,11 @@ Result> RawFileSplitRead::CreateReader( auto raw_readers = ObjectUtils::MoveVector>(std::move(raw_file_readers)); - auto concat_batch_reader = std::make_unique(std::move(raw_readers), pool_); + auto concat_batch_reader = + std::make_unique(std::move(raw_readers), arrow_pool_); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch_reader, ApplyPredicateFilterIfNeeded(std::move(concat_batch_reader), predicate)); - return std::make_unique(std::move(batch_reader), pool_); + return std::make_unique(std::move(batch_reader), arrow_pool_); } Result> RawFileSplitRead::CreateReader( diff --git a/src/paimon/core/operation/raw_file_split_read_test.cpp b/src/paimon/core/operation/raw_file_split_read_test.cpp index 51c478325..e6c2184e9 100644 --- a/src/paimon/core/operation/raw_file_split_read_test.cpp +++ b/src/paimon/core/operation/raw_file_split_read_test.cpp @@ -29,6 +29,7 @@ #include "paimon/common/data/binary_row.h" #include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/core_options.h" #include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/io/data_file_meta.h" @@ -137,13 +138,14 @@ class RawFileSplitReadTest : public ::testing::Test { "multi_partition_append_table"; ReadContextBuilder context_builder(path); context_builder.SetReadFieldNames(read_schema->field_names()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); - ASSERT_OK_AND_ASSIGN(auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + ASSERT_OK_AND_ASSIGN( + auto internal_context, + InternalReadContext::Create(read_context, table_schema, table_schema->Options(), + GetSharedArrowPool(read_context->GetMemoryPool()))); auto data_splits = PrepareDataSplits(); const auto& core_options = internal_context->GetCoreOptions(); auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); @@ -172,14 +174,16 @@ class RawFileSplitReadTest : public ::testing::Test { split_read->CreateReader(split)); batch_readers.emplace_back(std::move(reader)); } - auto batch_reader = std::make_unique(std::move(batch_readers), pool_); + auto batch_reader = + std::make_unique(std::move(batch_readers), arrow_pool_); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_TRUE(result_array->Equals(expected_array)); } private: std::shared_ptr pool_ = GetDefaultPool(); + std::shared_ptr arrow_pool_ = GetSharedArrowPool(pool_); }; // test simple, recall all columns with sequence in data @@ -383,13 +387,14 @@ TEST_F(RawFileSplitReadTest, TestEmptyPlan) { "/orc/multi_partition_append_table.db/" "multi_partition_append_table"; ReadContextBuilder context_builder(path); - ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); - ASSERT_OK_AND_ASSIGN(auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + ASSERT_OK_AND_ASSIGN( + auto internal_context, + InternalReadContext::Create(read_context, table_schema, table_schema->Options(), + GetSharedArrowPool(read_context->GetMemoryPool()))); const auto& core_options = internal_context->GetCoreOptions(); auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); ASSERT_OK_AND_ASSIGN(std::vector external_paths, @@ -422,8 +427,9 @@ TEST_F(RawFileSplitReadTest, TestEmptyPlan) { std::vector> batch_readers; ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, split_read->CreateReader(data_split)); batch_readers.push_back(std::move(reader)); - auto batch_reader = std::make_unique(std::move(batch_readers), pool_); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + auto batch_reader = std::make_unique(std::move(batch_readers), arrow_pool_); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_EQ(result_array, nullptr); } @@ -432,12 +438,13 @@ TEST_F(RawFileSplitReadTest, TestMatch) { "/orc/pk_table_with_total_buckets.db/pk_table_with_total_buckets"; ReadContextBuilder context_builder(path); context_builder.SetReadFieldNames({"f0", "f1", "f2", "f3"}); - ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr internal_context, + InternalReadContext::Create(read_context, table_schema, table_schema->Options(), + GetSharedArrowPool(read_context->GetMemoryPool()))); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(/*thread_count=*/2)); auto split_read = std::make_unique( diff --git a/src/paimon/core/postpone/postpone_bucket_writer_test.cpp b/src/paimon/core/postpone/postpone_bucket_writer_test.cpp index 30357327a..dd20eca48 100644 --- a/src/paimon/core/postpone/postpone_bucket_writer_test.cpp +++ b/src/paimon/core/postpone/postpone_bucket_writer_test.cpp @@ -107,7 +107,7 @@ class PostponeBucketWriterTest : public ::testing::Test, ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_TRUE(expected_array->Equals(result_array)) << result_array->ToString() << "\n != \n" << expected_array->ToString(); } diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp index 1136243e8..6dc9bbca2 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store.cpp @@ -34,6 +34,7 @@ #include "paimon/common/table/special_fields.h" #include "paimon/common/types/row_kind.h" #include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/projected_array.h" @@ -154,6 +155,7 @@ class ArrowRealtimeStore::CommitBatchReader : public BatchReader { auto c_array = std::make_unique(); auto c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*result, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return ReadBatch(std::move(c_array), std::move(c_schema)); } @@ -219,6 +221,7 @@ class ArrowRealtimeStore::QueryBatchReader : public BatchReader { auto c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*output, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); return ReadBatchWithBitmap(ReadBatch(std::move(c_array), std::move(c_schema)), std::move(candidate_rows)); } @@ -457,7 +460,7 @@ Result>> ArrowRealtimeStore::CreateQuer std::unique_ptr reader = std::make_unique( arrow_view.get(), offset_begin, read_schema, predicate_filter, std::move(statistics_mapping), arrow_pool_, memory_pool_); - reader = std::make_unique(std::move(reader), memory_pool_); + reader = std::make_unique(std::move(reader), arrow_pool_); readers.push_back(std::move(reader)); } return readers; diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index dff12b589..0edbf7953 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -44,7 +44,8 @@ Result> ArrowRealtimeStoreFactory::Create( arrow::ImportSchema(request.write_schema.get())); switch (request.mode) { case RealtimeStoreMode::APPEND_ONLY: { - std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); + std::shared_ptr arrow_pool = + GetSharedArrowPool(request.memory_pool); return std::make_shared(imported_schema, request.statistics_mode, request.memory_pool, arrow_pool); } diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 632d64e16..0733d9e93 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -31,6 +31,7 @@ #include "paimon/common/table/special_fields.h" #include "paimon/common/types/row_kind.h" #include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/scope_guard.h" @@ -69,7 +70,7 @@ RealtimeAppendOnlyWriter::RealtimeAppendOnlyWriter( const std::shared_ptr& file_writer, const std::shared_ptr& input_schema, int64_t next_offset, const std::shared_ptr& memory_pool) - : memory_pool_(memory_pool), + : arrow_pool_(GetSharedArrowPool(memory_pool)), realtime_store_(realtime_store), file_writer_(file_writer), input_schema_(input_schema), @@ -123,7 +124,7 @@ Status RealtimeAppendOnlyWriter::FlushSegment( const std::shared_ptr& segment) { PAIMON_ASSIGN_OR_RAISE(std::vector> readers, realtime_store_->CreateCommitReaders(segment)); - ConcatBatchReader reader(std::move(readers), memory_pool_); + ConcatBatchReader reader(std::move(readers), arrow_pool_); ScopeGuard reader_guard([&reader]() { reader.Close(); }); const OffsetRange offset_range = segment->GetOffsetRange(); int64_t emitted_rows = 0; diff --git a/src/paimon/core/realtime/realtime_append_only_writer.h b/src/paimon/core/realtime/realtime_append_only_writer.h index a6190d3b0..d588d0b83 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.h +++ b/src/paimon/core/realtime/realtime_append_only_writer.h @@ -31,6 +31,7 @@ struct ArrowSchema; namespace arrow { +class MemoryPool; class Schema; } // namespace arrow @@ -77,7 +78,7 @@ class RealtimeAppendOnlyWriter : public BatchWriter { Status FlushSegment(const std::shared_ptr& segment); - std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; std::shared_ptr realtime_store_; std::shared_ptr file_writer_; std::shared_ptr input_schema_; diff --git a/src/paimon/core/table/source/append_only_table_read.cpp b/src/paimon/core/table/source/append_only_table_read.cpp index 12b3823a7..b51ba99bd 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -51,7 +51,7 @@ AppendOnlyTableRead::AppendOnlyTableRead(const std::shared_ptr& context, const std::shared_ptr& memory_pool, const std::shared_ptr& executor) - : TableRead(memory_pool), context_(context) { + : context_(context) { const auto& core_options = context->GetCoreOptions(); if (core_options.DataEvolutionEnabled()) { // add data evolution first @@ -111,7 +111,8 @@ Result> AppendOnlyTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } - return std::make_unique(std::move(readers), GetMemoryPool()); + return std::make_unique(std::move(readers), + context_->GetArrowMemoryPool()); } Result> AppendOnlyTableRead::CreateRealtimeReader( @@ -174,9 +175,10 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( return Status::Invalid("append-only real-time store returned a null query reader"); } if (context_->EnablePredicateFilter() && context_->GetPredicate()) { - PAIMON_ASSIGN_OR_RAISE(memory_reader, PredicateBatchReader::Create( - std::move(memory_reader), - context_->GetPredicate(), GetMemoryPool())); + PAIMON_ASSIGN_OR_RAISE( + memory_reader, + PredicateBatchReader::Create(std::move(memory_reader), context_->GetPredicate(), + context_->GetArrowMemoryPool())); } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr realtime_reader, RealtimeReader::Create(memory.read_view, std::move(memory_reader))); @@ -187,7 +189,7 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); + std::make_unique(std::move(readers), context_->GetArrowMemoryPool()); readers_guard.Release(); return result; } @@ -254,7 +256,7 @@ Result> AppendOnlyTableRead::CreateCountReader( } return std::make_unique(splits, context_->GetCoreOptions().GetFileSystem(), - GetMemoryPool()); + context_->GetMemoryPool()); } } // namespace paimon diff --git a/src/paimon/core/table/source/fallback_data_split_test.cpp b/src/paimon/core/table/source/fallback_data_split_test.cpp index 5ae97e19e..e37a7e190 100644 --- a/src/paimon/core/table/source/fallback_data_split_test.cpp +++ b/src/paimon/core/table/source/fallback_data_split_test.cpp @@ -27,6 +27,7 @@ #include "gtest/gtest.h" #include "paimon/common/data/binary_row.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/file_source.h" @@ -46,25 +47,32 @@ namespace paimon::test { namespace { class TrackingTableRead : public TableRead { public: - explicit TrackingTableRead(const std::shared_ptr& pool) : TableRead(pool) {} - Result> CreateReader( const std::shared_ptr& split) override { last_split_ = split; return std::unique_ptr(); } + Result> CreateReader( + const std::vector>& splits) override { + if (splits.size() != 1) { + return Status::Invalid("tracking table read expects one split"); + } + return CreateReader(splits[0]); + } + std::shared_ptr last_split_; }; } // namespace TEST(FallbackTableReadTest, RoutesIndexedSplitToMainTable) { std::shared_ptr pool = GetDefaultPool(); - auto main_table = std::make_unique(pool); - auto fallback_table = std::make_unique(pool); + auto main_table = std::make_unique(); + auto fallback_table = std::make_unique(); TrackingTableRead* main_table_ptr = main_table.get(); TrackingTableRead* fallback_table_ptr = fallback_table.get(); - FallbackTableRead table_read(std::move(main_table), std::move(fallback_table), pool); + FallbackTableRead table_read(std::move(main_table), std::move(fallback_table), + GetSharedArrowPool(pool)); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/"", /*data_files=*/{}); diff --git a/src/paimon/core/table/source/fallback_table_read.cpp b/src/paimon/core/table/source/fallback_table_read.cpp index 5f5957609..b677ab038 100644 --- a/src/paimon/core/table/source/fallback_table_read.cpp +++ b/src/paimon/core/table/source/fallback_table_read.cpp @@ -19,6 +19,7 @@ #include "paimon/core/table/source/fallback_table_read.h" +#include "paimon/common/reader/concat_batch_reader.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/fallback_data_split.h" #include "paimon/global_index/indexed_split.h" @@ -26,6 +27,17 @@ #include "paimon/table/source/data_split.h" namespace paimon { +Result> FallbackTableRead::CreateReader( + const std::vector>& splits) { + std::vector> readers; + readers.reserve(splits.size()); + for (const std::shared_ptr& 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), arrow_pool_); +} + Result> FallbackTableRead::CreateReader( const std::shared_ptr& split) { auto fallback_data_split = std::dynamic_pointer_cast(split); diff --git a/src/paimon/core/table/source/fallback_table_read.h b/src/paimon/core/table/source/fallback_table_read.h index 91af53af5..09acf0f6f 100644 --- a/src/paimon/core/table/source/fallback_table_read.h +++ b/src/paimon/core/table/source/fallback_table_read.h @@ -26,24 +26,30 @@ #include "paimon/result.h" #include "paimon/table/source/table_read.h" +namespace arrow { +class MemoryPool; +} + namespace paimon { class DataSplit; -class MemoryPool; class FallbackTableRead : public TableRead { public: FallbackTableRead(std::unique_ptr main_table, std::unique_ptr fallback_table, - const std::shared_ptr& memory_pool) - : TableRead(memory_pool), - main_table_(std::move(main_table)), - fallback_table_(std::move(fallback_table)) {} + const std::shared_ptr& arrow_pool) + : main_table_(std::move(main_table)), + fallback_table_(std::move(fallback_table)), + arrow_pool_(arrow_pool) {} Result> CreateReader(const std::shared_ptr& split) override; + Result> CreateReader( + const std::vector>& splits) override; private: std::unique_ptr main_table_; std::unique_ptr fallback_table_; + std::shared_ptr arrow_pool_; }; } // namespace paimon diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index afb852a6b..21282dd64 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -119,9 +119,8 @@ KeyValueTableRead::KeyValueTableRead( const std::shared_ptr& path_factory, const std::shared_ptr& context, const std::shared_ptr& realtime_primary_key_transport_schema, - const std::shared_ptr& memory_pool, const std::shared_ptr& executor) - : TableRead(memory_pool), - split_reads_(std::move(split_reads)), + const std::shared_ptr& executor) + : split_reads_(std::move(split_reads)), path_factory_(path_factory), context_(context), realtime_primary_key_transport_schema_(realtime_primary_key_transport_schema), @@ -149,7 +148,7 @@ Result> KeyValueTableRead::Create( return std::unique_ptr( new KeyValueTableRead(std::move(split_reads), path_factory, context, - realtime_primary_key_transport_schema, memory_pool, executor)); + realtime_primary_key_transport_schema, executor)); } void KeyValueTableRead::ForceKeepDelete(bool force_keep_delete) { @@ -259,7 +258,8 @@ Result> KeyValueTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } - return std::make_unique(std::move(readers), GetMemoryPool()); + return std::make_unique(std::move(readers), + context_->GetArrowMemoryPool()); } Result> KeyValueTableRead::CreateRealtimeReader( @@ -294,7 +294,8 @@ Result> KeyValueTableRead::CreateRealtimeReader( std::vector> memory_readers, CreateMemoryReaders(realtime_split, memory, realtime_primary_key_transport_schema_, merge_read->GetKeySchema(), merge_read->GetValueSchema(), - merge_read->GetKeyComparator(), context_, GetMemoryPool())); + merge_read->GetKeyComparator(), context_, + context_->GetMemoryPool())); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, merge_read->CreateRealtimeReader(realtime_split->DiskSplits(), std::move(memory_readers))); @@ -327,9 +328,9 @@ Result> KeyValueTableRead::CreateCountReader( return Status::NotImplemented("CreateCountReader with force_keep_delete is not supported"); } - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr pk_count_reader, - PKCountReader::Create(splits, path_factory_, context_, GetMemoryPool(), executor_)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr pk_count_reader, + PKCountReader::Create(splits, path_factory_, context_, + context_->GetMemoryPool(), executor_)); return pk_count_reader; } diff --git a/src/paimon/core/table/source/key_value_table_read.h b/src/paimon/core/table/source/key_value_table_read.h index 1dd59b016..78017113b 100644 --- a/src/paimon/core/table/source/key_value_table_read.h +++ b/src/paimon/core/table/source/key_value_table_read.h @@ -60,7 +60,6 @@ class KeyValueTableRead : public TableRead { const std::shared_ptr& path_factory, const std::shared_ptr& context, const std::shared_ptr& realtime_primary_key_transport_schema, - const std::shared_ptr& memory_pool, const std::shared_ptr& executor); Result> CreateRealtimeReader( diff --git a/src/paimon/core/table/source/pk_count_reader_test.cpp b/src/paimon/core/table/source/pk_count_reader_test.cpp index 72507fee8..07f677afd 100644 --- a/src/paimon/core/table/source/pk_count_reader_test.cpp +++ b/src/paimon/core/table/source/pk_count_reader_test.cpp @@ -26,6 +26,7 @@ #include "gtest/gtest.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/table/source/data_split_impl.h" @@ -71,13 +72,15 @@ class PKCountReaderTest : public testing::Test { Result> CreateInternalContext( const std::string& table_path) { ReadContextBuilder read_context_builder(table_path); - PAIMON_ASSIGN_OR_RAISE(auto read_context, read_context_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, + read_context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), table_path); PAIMON_ASSIGN_OR_RAISE(auto table_schema, schema_manager.ReadSchema(0)); - PAIMON_ASSIGN_OR_RAISE(auto internal_context, - InternalReadContext::Create(std::move(read_context), table_schema, - table_schema->Options())); + PAIMON_ASSIGN_OR_RAISE( + auto internal_context, + InternalReadContext::Create(read_context, table_schema, table_schema->Options(), + GetSharedArrowPool(read_context->GetMemoryPool()))); return std::shared_ptr(std::move(internal_context)); } diff --git a/src/paimon/core/table/source/table_read.cpp b/src/paimon/core/table/source/table_read.cpp index 3d6a36766..ef6d80068 100644 --- a/src/paimon/core/table/source/table_read.cpp +++ b/src/paimon/core/table/source/table_read.cpp @@ -26,8 +26,8 @@ #include #include "fmt/format.h" -#include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/internal_read_context.h" @@ -52,7 +52,8 @@ class MemoryPool; namespace { Result> CreateInternalReadContext( - const std::shared_ptr& context, const std::string& branch) { + const std::shared_ptr& context, const std::string& branch, + const std::shared_ptr& arrow_pool) { std::map tmp_options = context->GetOptions(); std::shared_ptr table_schema; const auto& specific_table_schema = context->GetSpecificTableSchema(); @@ -82,7 +83,7 @@ Result> CreateInternalReadContext( if (branch != BranchManager::DEFAULT_MAIN_BRANCH) { options[Options::BRANCH] = branch; } - return InternalReadContext::Create(context, table_schema, options); + return InternalReadContext::Create(context, table_schema, options, arrow_pool); } Result> CreateTableRead( @@ -116,8 +117,9 @@ Result> CreateTableRead( } Result> NewDataTableRead(const std::shared_ptr& context) { + std::shared_ptr arrow_pool = GetSharedArrowPool(context->GetMemoryPool()); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_context, - CreateInternalReadContext(context, context->GetBranch())); + CreateInternalReadContext(context, context->GetBranch(), arrow_pool)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr table_read, @@ -132,19 +134,17 @@ Result> NewDataTableRead(const std::shared_ptr fallback_context, - CreateInternalReadContext(context, /*branch=*/scan_fallback_branch.value())); + CreateInternalReadContext(context, /*branch=*/scan_fallback_branch.value(), arrow_pool)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr fallback_table_read, CreateTableRead(fallback_context, context->GetMemoryPool(), context->GetExecutor())); - return std::make_unique( - std::move(table_read), std::move(fallback_table_read), context->GetMemoryPool()); + return std::make_unique(std::move(table_read), + std::move(fallback_table_read), arrow_pool); } } // namespace -TableRead::TableRead(const std::shared_ptr& memory_pool) : pool_(memory_pool) {} - Result> TableRead::Create(std::unique_ptr ctx) { std::shared_ptr context = std::move(ctx); if (context == nullptr) { @@ -173,17 +173,6 @@ Result> TableRead::Create(std::unique_ptr> TableRead::CreateReader( - const std::vector>& splits) { - std::vector> batch_readers; - batch_readers.reserve(splits.size()); - for (const auto& split : splits) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, CreateReader(split)); - batch_readers.emplace_back(std::move(reader)); - } - return std::make_unique(std::move(batch_readers), pool_); -} - Result> TableRead::CreateCountReader( const std::vector>& splits) { (void)splits; 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 7572f906e..7e35a39ca 100644 --- a/src/paimon/core/table/system/audit_log_system_table.cpp +++ b/src/paimon/core/table/system/audit_log_system_table.cpp @@ -69,14 +69,14 @@ class ChangelogBatchReader : public BatchReader { ChangelogBatchReader(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) + bool pack_update_before_after, + const std::shared_ptr& arrow_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()) {} + arrow_pool_(arrow_pool) {} Result NextBatch() override { while (true) { @@ -124,8 +124,8 @@ class ChangelogBatchReader : public BatchReader { 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, converter_->ConvertDataColumn( + array, row_group_lengths, arrow_pool_.get())); PAIMON_ASSIGN_OR_RAISE(array, CopyToStablePool(array)); output_arrays.push_back(array); } @@ -137,6 +137,7 @@ class ChangelogBatchReader : public BatchReader { 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())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(output_c_array.get(), arrow_pool_)); return std::make_pair(std::move(output_c_array), std::move(output_c_schema)); } } @@ -158,7 +159,7 @@ class ChangelogBatchReader : public BatchReader { } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::shared_ptr combined, - arrow::Concatenate({pending_update_before_, struct_array}, arrow_pool_)); + arrow::Concatenate({pending_update_before_, struct_array}, arrow_pool_.get())); pending_update_before_.reset(); if (!combined || combined->type_id() != arrow::Type::STRUCT) { return Status::Invalid("failed to concatenate binlog struct batches"); @@ -213,7 +214,7 @@ class ChangelogBatchReader : public BatchReader { /// The imported data batch may release its C Arrow buffers after this wrapper returns. /// Keep returned system-table arrays independent of that input batch lifetime. PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, - arrow::Concatenate({array}, arrow_pool_)); + arrow::Concatenate({array}, arrow_pool_.get())); return result; } @@ -226,7 +227,7 @@ class ChangelogBatchReader : public BatchReader { } std::shared_ptr value_kind_array = checked_pointer_cast(value_kind); - arrow::StringBuilder builder(arrow_pool_); + arrow::StringBuilder builder(arrow_pool_.get()); 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) { @@ -257,7 +258,7 @@ class ChangelogBatchReader : public BatchReader { } std::shared_ptr sequence_array = checked_pointer_cast(sequence); - arrow::Int64Builder builder(arrow_pool_); + arrow::Int64Builder builder(arrow_pool_.get()); 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) { @@ -278,8 +279,7 @@ class ChangelogBatchReader : public BatchReader { 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 arrow_pool_; std::shared_ptr pending_update_before_; }; @@ -289,11 +289,11 @@ class ChangelogTableRead : public TableRead { std::shared_ptr output_schema, bool include_sequence_number, std::shared_ptr converter, const std::shared_ptr& pool) - : TableRead(pool), - data_read_(std::move(data_read)), + : data_read_(std::move(data_read)), output_schema_(std::move(output_schema)), include_sequence_number_(include_sequence_number), - converter_(std::move(converter)) {} + converter_(std::move(converter)), + arrow_pool_(GetSharedArrowPool(pool)) {} Result> CreateReader( const std::vector>& splits) override { @@ -308,13 +308,13 @@ class ChangelogTableRead : public TableRead { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, CreateReader(split)); readers.push_back(std::move(reader)); } - return std::make_unique(std::move(readers), GetMemoryPool()); + return std::make_unique(std::move(readers), arrow_pool_); } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, data_read_->CreateReader(splits)); - return CreateChangelogBatchReader(std::move(reader), output_schema_, - include_sequence_number_, converter_, - /*pack_update_before_after=*/false, GetMemoryPool()); + return std::make_unique( + std::move(reader), output_schema_, include_sequence_number_, converter_, + /*pack_update_before_after=*/false, arrow_pool_); } Result> CreateReader( @@ -330,9 +330,9 @@ class ChangelogTableRead : public TableRead { } pack_update_before_after = data_split->IsStreaming(); } - return CreateChangelogBatchReader(std::move(reader), output_schema_, - include_sequence_number_, converter_, - pack_update_before_after, GetMemoryPool()); + return std::make_unique(std::move(reader), output_schema_, + include_sequence_number_, converter_, + pack_update_before_after, arrow_pool_); } private: @@ -340,6 +340,7 @@ class ChangelogTableRead : public TableRead { std::shared_ptr output_schema_; bool include_sequence_number_; std::shared_ptr converter_; + std::shared_ptr arrow_pool_; }; } // namespace @@ -348,9 +349,9 @@ 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); + return std::make_unique( + std::move(reader), std::move(output_schema), include_sequence_number, std::move(converter), + pack_update_before_after, GetSharedArrowPool(pool)); } AuditLogSystemTable::AuditLogSystemTable(std::shared_ptr fs, std::string table_path, diff --git a/src/paimon/core/table/system/in_memory_system_table.cpp b/src/paimon/core/table/system/in_memory_system_table.cpp index f4d1aa67e..c05a012bb 100644 --- a/src/paimon/core/table/system/in_memory_system_table.cpp +++ b/src/paimon/core/table/system/in_memory_system_table.cpp @@ -40,8 +40,8 @@ namespace { class InMemorySystemTableBatchReader : public BatchReader { public: InMemorySystemTableBatchReader(std::shared_ptr table, - const std::shared_ptr& pool) - : table_(std::move(table)), arrow_pool_(GetArrowPool(pool)) {} + const std::shared_ptr& arrow_pool) + : table_(std::move(table)), arrow_pool_(arrow_pool) {} Result NextBatch() override { if (emitted_) { @@ -54,7 +54,7 @@ class InMemorySystemTableBatchReader : public BatchReader { return BatchReader::MakeEofBatch(); } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr converter, - GenericRowToArrowArrayConverter::Create(schema, arrow_pool_.get())); + GenericRowToArrowArrayConverter::Create(schema, arrow_pool_)); return converter->NextBatch(rows); } @@ -68,7 +68,7 @@ class InMemorySystemTableBatchReader : public BatchReader { private: std::shared_ptr table_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; bool emitted_ = false; }; @@ -76,7 +76,7 @@ class InMemorySystemTableRead : public TableRead { public: InMemorySystemTableRead(std::shared_ptr table, const std::shared_ptr& memory_pool) - : TableRead(memory_pool), table_(std::move(table)) {} + : table_(std::move(table)), arrow_pool_(GetSharedArrowPool(memory_pool)) {} Result> CreateReader( const std::vector>& splits) override { @@ -88,7 +88,7 @@ class InMemorySystemTableRead : public TableRead { return Status::Invalid("unsupported split for ", table_->Name(), " system table"); } } - return std::make_unique(table_, GetMemoryPool()); + return std::make_unique(table_, arrow_pool_); } Result> CreateReader( @@ -99,6 +99,7 @@ class InMemorySystemTableRead : public TableRead { private: std::shared_ptr table_; + std::shared_ptr arrow_pool_; }; } // namespace diff --git a/src/paimon/core/table/system/system_table_test.cpp b/src/paimon/core/table/system/system_table_test.cpp index 7e5847dbf..41ac2d8b4 100644 --- a/src/paimon/core/table/system/system_table_test.cpp +++ b/src/paimon/core/table/system/system_table_test.cpp @@ -27,6 +27,7 @@ #include #include "arrow/api.h" +#include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" #include "paimon/core/schema/table_schema.h" @@ -121,7 +122,7 @@ TEST(SystemTableTest, TestStreamingBinlogPacksUpdateAcrossBatches) { /*pack_update_before_after=*/true, GetDefaultPool()); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - ReadResultCollector::CollectResult(reader.get())); + ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected_array = arrow::ipc::internal::json::ArrayFromJSON(actual->type(), R"([ ["+I", 10, ["a"], [1]], @@ -154,7 +155,7 @@ TEST(SystemTableTest, TestStreamingBinlogEmitsUnmatchedUpdateBefore) { /*pack_update_before_after=*/true, GetDefaultPool()); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - ReadResultCollector::CollectResult(reader.get())); + ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr expected_array = arrow::ipc::internal::json::ArrayFromJSON(actual->type(), R"([ ["-U", ["b"], [2]] @@ -165,6 +166,42 @@ TEST(SystemTableTest, TestStreamingBinlogEmitsUnmatchedUpdateBefore) { << "expected: " << expected->ToString() << "\nactual: " << actual->ToString(); } +TEST(SystemTableTest, TestChangelogBatchOutlivesReader) { + std::unique_ptr unique_pool = GetMemoryPool(); + std::shared_ptr pool = std::move(unique_pool); + std::weak_ptr weak_pool = pool; + std::shared_ptr input_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("pk", arrow::utf8()), + }); + std::shared_ptr input = + arrow::ipc::internal::json::ArrayFromJSON(input_type, R"([[0, "a"]])").ValueOrDie(); + std::shared_ptr output_schema = arrow::schema({ + arrow::field("rowkind", arrow::utf8(), /*nullable=*/false), + arrow::field("pk", arrow::list(arrow::utf8())), + }); + 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, pool); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_GT(pool->CurrentUsage(), 0); + reader.reset(); + pool.reset(); + ASSERT_FALSE(weak_pool.expired()); + + auto& [c_array, c_schema] = batch; + arrow::Result> imported = + arrow::ImportArray(c_array.get(), c_schema.get()); + ASSERT_TRUE(imported.ok()) << imported.status().ToString(); + std::shared_ptr output = std::move(imported).ValueOrDie(); + ASSERT_EQ(output->length(), 1); + output.reset(); + ASSERT_TRUE(weak_pool.expired()); +} + TEST(SystemTableTest, TestReadOptimizedSystemTableRegistration) { ASSERT_TRUE(SystemTableLoader::IsSupported(ReadOptimizedSystemTable::kName)); diff --git a/src/paimon/core/utils/manifest_meta_reader.cpp b/src/paimon/core/utils/manifest_meta_reader.cpp index 52faecaee..3fca1156d 100644 --- a/src/paimon/core/utils/manifest_meta_reader.cpp +++ b/src/paimon/core/utils/manifest_meta_reader.cpp @@ -42,8 +42,8 @@ class MemoryPool; ManifestMetaReader::ManifestMetaReader(std::unique_ptr&& reader, const std::shared_ptr& target_type, - const std::shared_ptr& pool) - : reader_(std::move(reader)), target_type_(target_type), pool_(GetArrowPool(pool)) {} + const std::shared_ptr& arrow_pool) + : reader_(std::move(reader)), target_type_(target_type), pool_(arrow_pool) {} Result ManifestMetaReader::NextBatch() { PAIMON_ASSIGN_OR_RAISE(ReadBatch src_result, reader_->NextBatch()); @@ -62,6 +62,7 @@ Result ManifestMetaReader::NextBatch() { PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*target_array, target_c_arrow_array.get(), target_c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(target_c_arrow_array.get(), pool_)); return std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema)); } diff --git a/src/paimon/core/utils/manifest_meta_reader.h b/src/paimon/core/utils/manifest_meta_reader.h index 89beb728c..354132cf3 100644 --- a/src/paimon/core/utils/manifest_meta_reader.h +++ b/src/paimon/core/utils/manifest_meta_reader.h @@ -42,7 +42,7 @@ class PAIMON_EXPORT ManifestMetaReader : public BatchReader { public: ManifestMetaReader(std::unique_ptr&& reader, const std::shared_ptr& target_type, - const std::shared_ptr& pool); + const std::shared_ptr& arrow_pool); ~ManifestMetaReader() override { DoClose(); @@ -68,7 +68,7 @@ class PAIMON_EXPORT ManifestMetaReader : public BatchReader { std::unique_ptr reader_; std::shared_ptr target_type_; - std::unique_ptr pool_; + std::shared_ptr pool_; }; } // namespace paimon diff --git a/src/paimon/core/utils/objects_file.h b/src/paimon/core/utils/objects_file.h index 43d782019..09274d5e3 100644 --- a/src/paimon/core/utils/objects_file.h +++ b/src/paimon/core/utils/objects_file.h @@ -30,6 +30,7 @@ #include "paimon/cache/cache.h" #include "paimon/common/data/columnar/columnar_row.h" #include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/path_util.h" @@ -83,6 +84,7 @@ class ObjectsFile { std::shared_ptr path_factory_; std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::unique_ptr> serializer_; std::shared_ptr writer_builder_; std::unique_ptr to_array_converter_; @@ -107,6 +109,7 @@ ObjectsFile::ObjectsFile(const std::shared_ptr& file_system, const std::shared_ptr& pool) : path_factory_(path_factory), pool_(pool), + arrow_pool_(GetSharedArrowPool(pool)), serializer_(std::move(serializer)), writer_builder_(std::move(writer_builder)), file_system_(file_system), @@ -187,7 +190,7 @@ Status ObjectsFile::ReadArrowBatches( PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch_reader, reader_builder_->Build(file_input_stream)); auto reader = std::make_unique(std::move(batch_reader), - serializer_->GetDataType(), pool_); + serializer_->GetDataType(), arrow_pool_); while (true) { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch arrow_array, reader->NextBatch()); auto& c_array = arrow_array.first; diff --git a/src/paimon/format/avro/avro_file_batch_reader.cpp b/src/paimon/format/avro/avro_file_batch_reader.cpp index f8b00c7c6..7f1048dd5 100644 --- a/src/paimon/format/avro/avro_file_batch_reader.cpp +++ b/src/paimon/format/avro/avro_file_batch_reader.cpp @@ -40,11 +40,11 @@ AvroFileBatchReader::AvroFileBatchReader(const std::shared_ptr& inp const std::shared_ptr<::arrow::DataType>& file_data_type, std::unique_ptr<::avro::DataFileReaderBase>&& reader, std::unique_ptr&& array_builder, - std::unique_ptr&& arrow_pool, + const std::shared_ptr& arrow_pool, int32_t batch_size, const std::shared_ptr& pool) : pool_(pool), - arrow_pool_(std::move(arrow_pool)), + arrow_pool_(arrow_pool), input_stream_(input_stream), file_data_type_(file_data_type), reader_(std::move(reader)), @@ -65,7 +65,7 @@ void AvroFileBatchReader::DoClose() { Result> AvroFileBatchReader::Create( const std::shared_ptr& input_stream, int32_t batch_size, - const std::shared_ptr& pool) { + const std::shared_ptr& pool, const std::shared_ptr& arrow_pool) { if (batch_size <= 0) { return Status::Invalid( fmt::format("invalid batch size {}, must be larger than 0", batch_size)); @@ -75,12 +75,11 @@ Result> AvroFileBatchReader::Create( const auto& avro_file_schema = reader->dataSchema(); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::DataType> file_data_type, AvroSchemaConverter::AvroSchemaToArrowDataType(avro_file_schema)); - auto arrow_pool = GetArrowPool(pool); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::unique_ptr array_builder, arrow::MakeBuilder(file_data_type, arrow_pool.get())); return std::unique_ptr( new AvroFileBatchReader(input_stream, file_data_type, std::move(reader), - std::move(array_builder), std::move(arrow_pool), batch_size, pool)); + std::move(array_builder), arrow_pool, batch_size, pool)); } Result> AvroFileBatchReader::CreateDataFileReader( @@ -134,6 +133,7 @@ Result AvroFileBatchReader::NextBatch() { std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); previous_batch_row_count_ = c_array->length; return make_pair(std::move(c_array), std::move(c_schema)); } catch (const ::avro::Exception& e) { diff --git a/src/paimon/format/avro/avro_file_batch_reader.h b/src/paimon/format/avro/avro_file_batch_reader.h index ccd1f7705..05bdf912f 100644 --- a/src/paimon/format/avro/avro_file_batch_reader.h +++ b/src/paimon/format/avro/avro_file_batch_reader.h @@ -38,7 +38,8 @@ class AvroFileBatchReader : public FileBatchReader { public: static Result> Create( const std::shared_ptr& input_stream, int32_t batch_size, - const std::shared_ptr& pool); + const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool); ~AvroFileBatchReader() override; @@ -92,13 +93,13 @@ class AvroFileBatchReader : public FileBatchReader { const std::shared_ptr<::arrow::DataType>& file_data_type, std::unique_ptr<::avro::DataFileReaderBase>&& reader, std::unique_ptr&& array_builder, - std::unique_ptr&& arrow_pool, int32_t batch_size, + const std::shared_ptr& arrow_pool, int32_t batch_size, const std::shared_ptr& pool); static constexpr size_t BUFFER_SIZE = 1024 * 1024; // 1M std::shared_ptr pool_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::shared_ptr input_stream_; std::shared_ptr<::arrow::DataType> file_data_type_; std::unique_ptr<::avro::DataFileReaderBase> reader_; diff --git a/src/paimon/format/avro/avro_file_batch_reader_test.cpp b/src/paimon/format/avro/avro_file_batch_reader_test.cpp index 0a0a858d1..476af7bde 100644 --- a/src/paimon/format/avro/avro_file_batch_reader_test.cpp +++ b/src/paimon/format/avro/avro_file_batch_reader_test.cpp @@ -70,15 +70,16 @@ class AvroFileBatchReaderTest : public ::testing::Test, public ::testing::WithPa ASSERT_OK(out->Close()); } - std::pair, std::shared_ptr> ReadData( - const std::string& file_path, int32_t read_batch_size) { + std::shared_ptr ReadData(const std::string& file_path, + int32_t read_batch_size) { EXPECT_OK_AND_ASSIGN(auto reader_builder, file_format_->CreateReaderBuilder(read_batch_size)); EXPECT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path)); EXPECT_OK_AND_ASSIGN(auto batch_reader, reader_builder->Build(in)); EXPECT_OK_AND_ASSIGN(auto result_array, ::paimon::test::ReadResultCollector::CollectResult( batch_reader.get())); - return std::make_pair(std::move(batch_reader), result_array); + EXPECT_TRUE(batch_reader->GetReaderMetrics()); + return result_array; } private: @@ -90,7 +91,7 @@ class AvroFileBatchReaderTest : public ::testing::Test, public ::testing::WithPa TEST_F(AvroFileBatchReaderTest, TestReadDataWithNull) { std::string path = paimon::test::GetDataDir() + "/avro/data/avro_with_null"; - auto [reader_holder, result_array] = ReadData(path, /*read_batch_size=*/1024); + auto result_array = ReadData(path, /*read_batch_size=*/1024); arrow::FieldVector fields = { arrow::field("_KEY_f0", arrow::utf8(), /*nullable=*/true), @@ -113,8 +114,6 @@ TEST_F(AvroFileBatchReaderTest, TestReadDataWithNull) { ASSERT_TRUE(array_status.ok()) << array_status.ToString(); ASSERT_TRUE(result_array->Equals(expected_array)); ASSERT_TRUE(expected_array->Equals(result_array)); - auto read_metrics = reader_holder->GetReaderMetrics(); - ASSERT_TRUE(read_metrics); } TEST_F(AvroFileBatchReaderTest, TestReadWithDifferentBatchSize) { @@ -152,7 +151,7 @@ TEST_F(AvroFileBatchReaderTest, TestReadWithDifferentBatchSize) { WriteData(src_array, file_path, /*compression=*/"zstd"); for (int32_t batch_size : {1024, 512, 256, 128, 64, 32, 16, 8, 4, 2, 1}) { - auto [reader_holder, result_array] = ReadData(file_path, batch_size); + auto result_array = ReadData(file_path, batch_size); std::shared_ptr expected_array; auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON( arrow_data_type, {data_str}, &expected_array); @@ -164,7 +163,7 @@ TEST_F(AvroFileBatchReaderTest, TestReadWithDifferentBatchSize) { TEST_F(AvroFileBatchReaderTest, TestReadAllTypes) { std::string path = paimon::test::GetDataDir() + "/avro/data/avro_all_types"; - auto [reader_holder, result_array] = ReadData(path, /*read_batch_size=*/1024); + auto result_array = ReadData(path, /*read_batch_size=*/1024); arrow::FieldVector fields = { arrow::field("f0", arrow::boolean()), @@ -229,8 +228,8 @@ TEST_P(AvroFileBatchReaderTest, TestReadTimestampTypes) { /*selection_bitmap=*/std::nullopt)); // check array - ASSERT_OK_AND_ASSIGN(auto result_array, - ::paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ::paimon::test::ReadResultCollector::CollectResult( + std::move(batch_reader))); std::shared_ptr expected_array; auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(read_fields), {R"([ @@ -275,8 +274,8 @@ TEST_F(AvroFileBatchReaderTest, TestReadMapTypes) { /*selection_bitmap=*/std::nullopt)); // check array - ASSERT_OK_AND_ASSIGN(auto result_array, - ::paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ::paimon::test::ReadResultCollector::CollectResult( + std::move(batch_reader))); std::shared_ptr expected_array; auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(read_fields), {R"([ @@ -523,7 +522,7 @@ TEST_F(AvroFileBatchReaderTest, TestReadBinaryWrittenFromBinaryAndLargeBinary) { /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto result_array, ::paimon::test::ReadResultCollector::CollectResult( - batch_reader.get())); + std::move(batch_reader))); auto expected_array = arrow::ipc::internal::json::ArrayFromJSON(read_data_type, data_json).ValueOrDie(); auto expected_chunked_array = std::make_shared(expected_array); diff --git a/src/paimon/format/avro/avro_file_format_test.cpp b/src/paimon/format/avro/avro_file_format_test.cpp index 6008cdd0b..fe361940a 100644 --- a/src/paimon/format/avro/avro_file_format_test.cpp +++ b/src/paimon/format/avro/avro_file_format_test.cpp @@ -90,7 +90,7 @@ class AvroFileFormatTest : public testing::Test, public ::testing::WithParamInte ASSERT_OK(batch_reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto output_array, ::paimon::test::ReadResultCollector::CollectResult( - batch_reader.get())); + std::move(batch_reader))); ASSERT_TRUE(output_array->Equals(arrow::ChunkedArray(input_array))) << output_array->ToString() << "\n vs \n" << input_array->ToString(); diff --git a/src/paimon/format/avro/avro_format_writer_test.cpp b/src/paimon/format/avro/avro_format_writer_test.cpp index 1e6e84e17..f89d6b7bb 100644 --- a/src/paimon/format/avro/avro_format_writer_test.cpp +++ b/src/paimon/format/avro/avro_format_writer_test.cpp @@ -121,12 +121,12 @@ class AvroFormatWriterTest : public ::testing::Test { void CheckResult(const std::string& file_path, int32_t row_count) const { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs_->Open(file_path)); ASSERT_OK_AND_ASSIGN(auto file_reader, - AvroFileBatchReader::Create(input_stream, 1024, pool_)); + AvroFileBatchReader::Create(input_stream, 1024, pool_, arrow_pool_)); ASSERT_OK_AND_ASSIGN(uint64_t num_rows, file_reader->GetNumberOfRows()); ASSERT_EQ(num_rows, row_count); - ASSERT_OK_AND_ASSIGN(auto result_array, - ::paimon::test::ReadResultCollector::CollectResult(file_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, ::paimon::test::ReadResultCollector::CollectResult( + std::move(file_reader))); const auto& struct_array = checked_pointer_cast(result_array->chunk(0)); const auto& string_array = checked_pointer_cast(struct_array->field(0)); ASSERT_TRUE(string_array); diff --git a/src/paimon/format/avro/avro_reader_builder.h b/src/paimon/format/avro/avro_reader_builder.h index f19d86834..4d450d4ae 100644 --- a/src/paimon/format/avro/avro_reader_builder.h +++ b/src/paimon/format/avro/avro_reader_builder.h @@ -21,6 +21,7 @@ #include #include "avro/DataFile.hh" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/avro/avro_file_batch_reader.h" #include "paimon/format/avro/avro_input_stream_impl.h" #include "paimon/format/reader_builder.h" @@ -32,21 +33,25 @@ namespace paimon::avro { class AvroReaderBuilder : public ReaderBuilder { public: explicit AvroReaderBuilder(int32_t batch_size) - : batch_size_(batch_size), pool_(GetDefaultPool()) {} + : batch_size_(batch_size), + pool_(GetDefaultPool()), + arrow_pool_(GetSharedArrowPool(pool_)) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; + arrow_pool_ = GetSharedArrowPool(pool); return this; } Result> Build( const std::shared_ptr& path) const override { - return AvroFileBatchReader::Create(path, batch_size_, pool_); + return AvroFileBatchReader::Create(path, batch_size_, pool_, arrow_pool_); } private: const int32_t batch_size_; std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; } // namespace paimon::avro diff --git a/src/paimon/format/blob/blob_file_batch_reader.cpp b/src/paimon/format/blob/blob_file_batch_reader.cpp index d52ac8ead..8f7b0649d 100644 --- a/src/paimon/format/blob/blob_file_batch_reader.cpp +++ b/src/paimon/format/blob/blob_file_batch_reader.cpp @@ -42,7 +42,8 @@ namespace paimon::blob { Result> BlobFileBatchReader::Create( const std::shared_ptr& input_stream, int32_t batch_size, bool blob_as_descriptor, - bool emit_placeholder_sentinel, const std::shared_ptr& pool) { + bool emit_placeholder_sentinel, const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool) { if (input_stream == nullptr) { return Status::Invalid("blob file batch reader create failed: input stream is nullptr"); } @@ -95,17 +96,15 @@ Result> BlobFileBatchReader::Create( PAIMON_ASSIGN_OR_RAISE(std::string file_path, input_stream->GetUri()); auto reader = std::unique_ptr( new BlobFileBatchReader(input_stream, file_path, blob_lengths, blob_offsets, batch_size, - blob_as_descriptor, emit_placeholder_sentinel, pool)); + blob_as_descriptor, emit_placeholder_sentinel, pool, arrow_pool)); return reader; } -BlobFileBatchReader::BlobFileBatchReader(const std::shared_ptr& input_stream, - const std::string& file_path, - const std::vector& blob_lengths, - const std::vector& blob_offsets, - int32_t batch_size, bool blob_as_descriptor, - bool emit_placeholder_sentinel, - const std::shared_ptr& pool) +BlobFileBatchReader::BlobFileBatchReader( + const std::shared_ptr& input_stream, const std::string& file_path, + const std::vector& blob_lengths, const std::vector& blob_offsets, + int32_t batch_size, bool blob_as_descriptor, bool emit_placeholder_sentinel, + const std::shared_ptr& pool, const std::shared_ptr& arrow_pool) : input_stream_(input_stream), file_path_(file_path), all_blob_lengths_(blob_lengths), @@ -116,7 +115,7 @@ BlobFileBatchReader::BlobFileBatchReader(const std::shared_ptr& inp blob_as_descriptor_(blob_as_descriptor), emit_placeholder_sentinel_(emit_placeholder_sentinel), pool_(pool), - arrow_pool_(GetArrowPool(pool_)), + arrow_pool_(arrow_pool), metrics_(std::make_shared()) { target_blob_row_indexes_.resize(target_blob_lengths_.size()); std::iota(target_blob_row_indexes_.begin(), target_blob_row_indexes_.end(), 0); @@ -326,6 +325,7 @@ Result BlobFileBatchReader::NextBatch() { std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*blob_array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); previous_batch_start_pos_ = current_pos_; current_pos_ += rows_to_read; previous_batch_row_count_ = c_array->length; diff --git a/src/paimon/format/blob/blob_file_batch_reader.h b/src/paimon/format/blob/blob_file_batch_reader.h index ccb0afe80..9a05c15b1 100644 --- a/src/paimon/format/blob/blob_file_batch_reader.h +++ b/src/paimon/format/blob/blob_file_batch_reader.h @@ -105,7 +105,8 @@ class BlobFileBatchReader : public FileBatchReader { static Result> Create( const std::shared_ptr& input_stream, int32_t batch_size, bool blob_as_descriptor, bool emit_placeholder_sentinel, - const std::shared_ptr& pool); + const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool); Result> GetFileSchema() const override; @@ -157,7 +158,8 @@ class BlobFileBatchReader : public FileBatchReader { const std::string& file_path, const std::vector& blob_lengths, const std::vector& blob_offsets, int32_t batch_size, bool blob_as_descriptor, bool emit_placeholder_sentinel, - const std::shared_ptr& pool); + const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool); Status ReadBlobContentAt(const int64_t offset, const int64_t length, uint8_t* content) const; diff --git a/src/paimon/format/blob/blob_file_batch_reader_test.cpp b/src/paimon/format/blob/blob_file_batch_reader_test.cpp index 95637720a..f0a200bda 100644 --- a/src/paimon/format/blob/blob_file_batch_reader_test.cpp +++ b/src/paimon/format/blob/blob_file_batch_reader_test.cpp @@ -22,6 +22,7 @@ #include "arrow/c/helpers.h" #include "gtest/gtest.h" #include "paimon/common/data/blob_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/data/blob.h" #include "paimon/format/blob/blob_format_writer.h" #include "paimon/fs/local/local_file_system.h" @@ -35,6 +36,7 @@ class BlobFileBatchReaderTest : public testing::Test, public ::testing::WithPara public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); } void CheckResult(const std::string& table_path, const std::string& paimon_blob_file, @@ -46,12 +48,13 @@ class BlobFileBatchReaderTest : public testing::Test, public ::testing::WithPara std::shared_ptr fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/" + paimon_blob_file)); - ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( - input_stream, /*batch_size=*/1024, blob_as_descriptor, - /*emit_placeholder_sentinel=*/false, pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, selection_bitmap)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); if (chunked_array == nullptr) { ASSERT_EQ(0, original_blob_files.size()); return; @@ -102,6 +105,7 @@ class BlobFileBatchReaderTest : public testing::Test, public ::testing::WithPara private: std::string blob_field_name_; std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; TEST_P(BlobFileBatchReaderTest, TestSimple) { @@ -164,10 +168,10 @@ TEST_F(BlobFileBatchReaderTest, TestRowNumbers) { ASSERT_OK_AND_ASSIGN( std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); - ASSERT_OK_AND_ASSIGN(auto reader, - BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( + input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, std::nullopt)); ASSERT_OK_AND_ASSIGN(auto number_of_rows, reader->GetNumberOfRows()); @@ -206,10 +210,10 @@ TEST_F(BlobFileBatchReaderTest, TestRowNumbersWithSelectionBitmap) { ASSERT_OK_AND_ASSIGN( std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); - ASSERT_OK_AND_ASSIGN(auto reader, - BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( + input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); RoaringBitmap32 selection; selection.Add(0); @@ -245,21 +249,22 @@ TEST_F(BlobFileBatchReaderTest, InvalidScenario) { ASSERT_NOK_WITH_MSG( BlobFileBatchReader::Create(input_stream, /*batch_size=*/0, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_), + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_), "blob file batch reader create failed: read batch size '0' should be larger than zero"); } { ASSERT_NOK_WITH_MSG( BlobFileBatchReader::Create(/*input_stream=*/nullptr, /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_), + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_), "blob file batch reader create failed: input stream is nullptr"); } { ASSERT_OK_AND_ASSIGN( - auto reader, BlobFileBatchReader::Create(/*input_stream=*/input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_)); + auto reader, + BlobFileBatchReader::Create(/*input_stream=*/input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); ASSERT_NOK_WITH_MSG(reader->GetFileSchema(), "blob file has no self-describing file schema"); ASSERT_TRUE(reader->GetReaderMetrics()); @@ -293,10 +298,10 @@ TEST_P(BlobFileBatchReaderTest, EmptyFile) { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system->Open(dir->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN(auto reader, - BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( + input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, std::nullopt)); ASSERT_OK_AND_ASSIGN(auto number_of_rows, reader->GetNumberOfRows()); @@ -319,9 +324,10 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); ASSERT_OK_AND_ASSIGN( - auto reader, BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_)); + auto reader, + BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(/*read_schema=*/nullptr, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "SetReadSchema failed: read schema cannot be nullptr"); @@ -344,9 +350,10 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); ASSERT_OK_AND_ASSIGN( - auto reader, BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_)); + auto reader, + BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "read schema field number 2 is not 1"); @@ -369,9 +376,10 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); ASSERT_OK_AND_ASSIGN( - auto reader, BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_)); + auto reader, + BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "field my_blob_field: large_binary is not BLOB"); @@ -392,9 +400,10 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); ASSERT_OK_AND_ASSIGN( - auto reader, BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_)); + auto reader, + BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); RoaringBitmap32 roaring; roaring.Add(0); roaring.Add(1); diff --git a/src/paimon/format/blob/blob_format_writer_test.cpp b/src/paimon/format/blob/blob_format_writer_test.cpp index a813d66a1..b0afc7ace 100644 --- a/src/paimon/format/blob/blob_format_writer_test.cpp +++ b/src/paimon/format/blob/blob_format_writer_test.cpp @@ -25,6 +25,7 @@ #include "arrow/c/bridge.h" #include "gtest/gtest.h" #include "paimon/common/data/blob_descriptor.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/stream_utils.h" @@ -115,6 +116,7 @@ class BlobFormatWriterTestBase : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + arrow_pool_ = GetSharedArrowPool(pool_); dir_ = paimon::test::UniqueTestDirectory::Create(); ASSERT_TRUE(dir_); file_system_ = std::make_shared(); @@ -176,14 +178,14 @@ class BlobFormatWriterTestBase : public ::testing::Test { std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/false, pool_)); + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); PAIMON_RETURN_NOT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr concat_array, arrow::Concatenate(chunked_array->chunks())); return checked_pointer_cast(concat_array); @@ -191,6 +193,7 @@ class BlobFormatWriterTestBase : public ::testing::Test { protected: std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::unique_ptr dir_; std::shared_ptr output_stream_; std::shared_ptr file_system_; @@ -267,14 +270,14 @@ TEST_P(BlobFormatWriterTest, TestSimple) { ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_)); + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK( reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); // check result if (blob_as_descriptor_) { @@ -464,14 +467,14 @@ TEST_P(BlobFormatWriterTest, TestLargeBlob) { ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_)); + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK( reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); // check result if (blob_as_descriptor_) { @@ -514,14 +517,14 @@ TEST_P(BlobFormatWriterTest, TestAddBatchWithNullValues) { ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_)); + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK( reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); auto concat_array = arrow::Concatenate(chunked_array->chunks()).ValueOrDie(); auto result_struct = checked_pointer_cast(concat_array); @@ -1042,7 +1045,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/false, pool_)); + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, @@ -1058,13 +1061,13 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/true, pool_)); + /*emit_placeholder_sentinel=*/true, pool_, arrow_pool_)); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); auto concat_array = arrow::Concatenate(chunked_array->chunks()).ValueOrDie(); auto struct_array = checked_pointer_cast(concat_array); ASSERT_EQ(struct_array->length(), 2); @@ -1082,13 +1085,13 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/true, pool_)); + /*emit_placeholder_sentinel=*/true, pool_, arrow_pool_)); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); auto concat_array = arrow::Concatenate(chunked_array->chunks()).ValueOrDie(); auto struct_array = checked_pointer_cast(concat_array); auto binary_array = checked_pointer_cast(struct_array->field(0)); @@ -1109,10 +1112,11 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderWithSelectionBitmap) ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/true, pool_)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/false, + /*emit_placeholder_sentinel=*/true, pool_, arrow_pool_)); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -1121,7 +1125,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderWithSelectionBitmap) selection.Add(2); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, selection)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); auto concat_array = arrow::Concatenate(chunked_array->chunks()).ValueOrDie(); auto struct_array = checked_pointer_cast(concat_array); ASSERT_EQ(struct_array->length(), 2); @@ -1141,17 +1145,18 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestSentinelBytesVerbatimWithoutPlacehol ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/false, pool_)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/false, + /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); auto concat_array = arrow::Concatenate(chunked_array->chunks()).ValueOrDie(); auto struct_array = checked_pointer_cast(concat_array); ASSERT_EQ(struct_array->length(), 1); @@ -1177,16 +1182,17 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestSentinelPrefixedValueVerbatimInPlace for (bool emit_placeholder_sentinel : {false, true}) { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/false, - emit_placeholder_sentinel, pool_)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/false, emit_placeholder_sentinel, + pool_, arrow_pool_)); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto chunked_array, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); auto concat_array = arrow::Concatenate(chunked_array->chunks()).ValueOrDie(); auto struct_array = checked_pointer_cast(concat_array); auto binary_array = checked_pointer_cast(struct_array->field(0)); diff --git a/src/paimon/format/blob/blob_reader_builder.h b/src/paimon/format/blob/blob_reader_builder.h index 46bb1d251..c66465922 100644 --- a/src/paimon/format/blob/blob_reader_builder.h +++ b/src/paimon/format/blob/blob_reader_builder.h @@ -22,6 +22,7 @@ #include #include +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/options_utils.h" #include "paimon/format/blob/blob_file_batch_reader.h" #include "paimon/format/reader_builder.h" @@ -33,10 +34,14 @@ namespace paimon::blob { class BlobReaderBuilder : public ReaderBuilder { public: BlobReaderBuilder(int32_t batch_size, const std::map& options) - : batch_size_(batch_size), pool_(GetDefaultPool()), options_(options) {} + : batch_size_(batch_size), + pool_(GetDefaultPool()), + arrow_pool_(GetSharedArrowPool(pool_)), + options_(options) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; + arrow_pool_ = GetSharedArrowPool(pool); return this; } @@ -49,12 +54,13 @@ class BlobReaderBuilder : public ReaderBuilder { OptionsUtils::GetValueFromMap( options_, BlobDefs::kEmitPlaceholderSentinelKey, false)); return BlobFileBatchReader::Create(input_stream, batch_size_, blob_as_descriptor, - emit_placeholder_sentinel, pool_); + emit_placeholder_sentinel, pool_, arrow_pool_); } private: int32_t batch_size_; std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::map options_; }; diff --git a/src/paimon/format/blob/blob_stats_extractor.cpp b/src/paimon/format/blob/blob_stats_extractor.cpp index ca671e4e8..8cbbafe09 100644 --- a/src/paimon/format/blob/blob_stats_extractor.cpp +++ b/src/paimon/format/blob/blob_stats_extractor.cpp @@ -24,6 +24,7 @@ #include "arrow/api.h" #include "fmt/format.h" #include "paimon/common/data/blob_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/blob/blob_file_batch_reader.h" #include "paimon/format/column_stats.h" #include "paimon/fs/file_system.h" @@ -56,7 +57,8 @@ BlobStatsExtractor::ExtractWithFileInfo(const std::shared_ptr& file_ std::unique_ptr blob_reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool)); + /*emit_placeholder_sentinel=*/false, pool, + GetSharedArrowPool(pool))); ColumnStatsVector result_stats; result_stats.push_back( ColumnStats::CreateStringColumnStats(std::nullopt, std::nullopt, std::nullopt)); diff --git a/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp b/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp index b065c3632..097b04584 100644 --- a/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp +++ b/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp @@ -58,12 +58,11 @@ MosaicFileBatchReader::MosaicFileBatchReader( Result> MosaicFileBatchReader::Create( const std::shared_ptr& input, int32_t batch_size, - const std::shared_ptr& pool) { + const std::shared_ptr& pool, const std::shared_ptr& arrow_pool) { if (input == nullptr || pool == nullptr || batch_size <= 0) { return Status::Invalid( "Mosaic reader requires non-null input and memory pool, and positive batch size"); } - std::shared_ptr arrow_pool = GetArrowPool(pool); PAIMON_ASSIGN_OR_RAISE(int64_t signed_length, input->Length()); PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(signed_length, "Mosaic input length")); auto length = static_cast(signed_length); @@ -200,6 +199,7 @@ Result MosaicFileBatchReader::NextBatch() { auto ffi_schema = std::make_unique<::ArrowSchema>(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*normalized_array, ffi_array.get(), ffi_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(ffi_array.get(), arrow_pool_)); return std::make_pair(std::move(ffi_array), std::move(ffi_schema)); } diff --git a/src/paimon/format/mosaic/mosaic_file_batch_reader.h b/src/paimon/format/mosaic/mosaic_file_batch_reader.h index 030e26c44..79ed2c691 100644 --- a/src/paimon/format/mosaic/mosaic_file_batch_reader.h +++ b/src/paimon/format/mosaic/mosaic_file_batch_reader.h @@ -42,7 +42,8 @@ class MosaicFileBatchReader : public FileBatchReader { public: static Result> Create( const std::shared_ptr& input, int32_t batch_size, - const std::shared_ptr& pool); + const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool); ~MosaicFileBatchReader() override; diff --git a/src/paimon/format/mosaic/mosaic_file_format_test.cpp b/src/paimon/format/mosaic/mosaic_file_format_test.cpp index 5dc37b5af..d356fb26b 100644 --- a/src/paimon/format/mosaic/mosaic_file_format_test.cpp +++ b/src/paimon/format/mosaic/mosaic_file_format_test.cpp @@ -420,7 +420,7 @@ TEST_F(MosaicFileFormatTest, RowGroupPredicateFiltering) { PredicateBuilder::IsNull(/*field_index=*/1, /*field_name=*/"untracked", FieldType::INT); ASSERT_OK(reader->SetReadSchema(&ffi_schema, predicate, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual_is_null_without_stats, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); ASSERT_TRUE(actual_is_null_without_stats->Equals(arrow::ChunkedArray(data))) << actual_is_null_without_stats->ToString(); } diff --git a/src/paimon/format/mosaic/mosaic_reader_builder.h b/src/paimon/format/mosaic/mosaic_reader_builder.h index ac1afba92..34bc11eb9 100644 --- a/src/paimon/format/mosaic/mosaic_reader_builder.h +++ b/src/paimon/format/mosaic/mosaic_reader_builder.h @@ -21,6 +21,7 @@ #include +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/mosaic/mosaic_file_batch_reader.h" #include "paimon/format/reader_builder.h" #include "paimon/memory/memory_pool.h" @@ -30,10 +31,13 @@ namespace paimon::mosaic { class MosaicReaderBuilder : public ReaderBuilder { public: explicit MosaicReaderBuilder(int32_t batch_size) - : batch_size_(batch_size), pool_(GetDefaultPool()) {} + : batch_size_(batch_size), + pool_(GetDefaultPool()), + arrow_pool_(GetSharedArrowPool(pool_)) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; + arrow_pool_ = GetSharedArrowPool(pool); return this; } @@ -42,12 +46,13 @@ class MosaicReaderBuilder : public ReaderBuilder { if (pool_ == nullptr) { return Status::Invalid("Mosaic reader memory pool is nullptr"); } - return MosaicFileBatchReader::Create(input, batch_size_, pool_); + return MosaicFileBatchReader::Create(input, batch_size_, pool_, arrow_pool_); } private: int32_t batch_size_; std::shared_ptr pool_; + std::shared_ptr arrow_pool_; }; } // namespace paimon::mosaic diff --git a/src/paimon/format/orc/complex_predicate_test.cpp b/src/paimon/format/orc/complex_predicate_test.cpp index 079cb319f..e7f8e1fc6 100644 --- a/src/paimon/format/orc/complex_predicate_test.cpp +++ b/src/paimon/format/orc/complex_predicate_test.cpp @@ -63,6 +63,7 @@ class ComplexPredicateTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + read_memory_ = std::make_shared(pool_); batch_size_ = 10; } void TearDown() override {} @@ -79,7 +80,7 @@ class ComplexPredicateTest : public ::testing::Test { EXPECT_TRUE(in_stream); EXPECT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), pool_, + OrcFileBatchReader::Create(std::move(in_stream), read_memory_, /*options=*/{{"orc.timestamp-ltz.legacy.type", "false"}}, batch_size)); EXPECT_TRUE(orc_batch_reader); @@ -98,7 +99,7 @@ class ComplexPredicateTest : public ::testing::Test { auto orc_batch_reader = PrepareOrcFileBatchReader(file_name, read_schema.get(), predicate, batch_size_); ASSERT_OK_AND_ASSIGN(auto arrow_array, paimon::test::ReadResultCollector::CollectResult( - orc_batch_reader.get())); + std::move(orc_batch_reader))); // check result if (expected_array) { ASSERT_TRUE(arrow_array); @@ -111,6 +112,7 @@ class ComplexPredicateTest : public ::testing::Test { private: std::shared_ptr pool_; + std::shared_ptr read_memory_; int32_t batch_size_; }; diff --git a/src/paimon/format/orc/orc_adapter_test.cpp b/src/paimon/format/orc/orc_adapter_test.cpp index 1a43de5ff..a59feee9a 100644 --- a/src/paimon/format/orc/orc_adapter_test.cpp +++ b/src/paimon/format/orc/orc_adapter_test.cpp @@ -60,8 +60,8 @@ class OrcAdapterTest : public ::testing::Test, void SetUp() override {} void TearDown() override {} - std::pair, std::unique_ptr<::orc::ColumnVectorBatch>> - GenerateOrcReadBatch(const std::shared_ptr& src_array) const { + std::shared_ptr<::orc::ColumnVectorBatch> GenerateOrcReadBatch( + const std::shared_ptr& src_array) const { auto [dict_key_size_threshold, enable_lazy_decoding] = GetParam(); arrow::Schema src_schema(src_array->type()->fields()); EXPECT_OK_AND_ASSIGN(std::unique_ptr<::orc::Type> orc_type, @@ -107,7 +107,12 @@ class OrcAdapterTest : public ::testing::Test, std::unique_ptr<::orc::RowReader> row_reader = reader->createRowReader(options); auto read_batch = row_reader->createRowBatch(src_array->length() + 10); [[maybe_unused]] bool not_eof = row_reader->next(*read_batch); - return std::make_pair(std::move(row_reader), std::move(read_batch)); + std::shared_ptr<::orc::RowReader> row_reader_lifetime(std::move(row_reader)); + return std::shared_ptr<::orc::ColumnVectorBatch>( + read_batch.release(), [row_reader_lifetime](::orc::ColumnVectorBatch* batch) { + delete batch; + (void)row_reader_lifetime; + }); } }; @@ -385,7 +390,7 @@ TEST_P(OrcAdapterTest, TestEmptyBatch) { arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ ])") .ValueOrDie()); - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); ASSERT_OK_AND_ASSIGN( std::shared_ptr target_array, @@ -419,7 +424,7 @@ TEST_P(OrcAdapterTest, TestDictionary) { // test with dict (f0), without dict (f1) orthogonal with enable_lazy_decoding = {true, false} // all 3 conditions are touched - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); ASSERT_OK_AND_ASSIGN( std::shared_ptr target_array, OrcAdapter::AppendBatch(arrow_type, read_batch.get(), arrow::default_memory_pool())); @@ -449,7 +454,7 @@ TEST_P(OrcAdapterTest, TestShadowCopyWithBlob) { ["data", 4.0, 8.0, 70, 500, 50000, 40000, true, "data", true] ])") .ValueOrDie()); - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); ASSERT_OK_AND_ASSIGN( std::shared_ptr target_array, @@ -477,7 +482,7 @@ TEST_P(OrcAdapterTest, TestDeepCopyWithString) { ["data", "data"] ])") .ValueOrDie()); - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); ASSERT_OK_AND_ASSIGN( std::shared_ptr target_array, @@ -514,7 +519,7 @@ TEST_P(OrcAdapterTest, TestComplexTypeShallowCopyWithBlob) { ])") .ValueOrDie()); - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); ASSERT_OK_AND_ASSIGN( std::shared_ptr target_array, OrcAdapter::AppendBatch(arrow_type, read_batch.get(), arrow::default_memory_pool())); @@ -536,7 +541,7 @@ TEST_P(OrcAdapterTest, TestAppendBatchWithBinary) { ["data", "data"] ])") .ValueOrDie()); - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); ASSERT_OK_AND_ASSIGN( std::shared_ptr target_array, @@ -556,7 +561,7 @@ TEST_P(OrcAdapterTest, TestAppendBatchWithBinaryForAllNull) { [null] ])") .ValueOrDie()); - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); ASSERT_OK_AND_ASSIGN( std::shared_ptr target_array, @@ -578,7 +583,7 @@ TEST_P(OrcAdapterTest, TestWriteBatchWithLargeBinary) { ])") .ValueOrDie()); - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); auto* struct_batch = dynamic_cast<::orc::StructVectorBatch*>(read_batch.get()); ASSERT_TRUE(struct_batch); ASSERT_EQ(1, struct_batch->fields.size()); @@ -632,7 +637,7 @@ TEST_P(OrcAdapterTest, TestDecimalAndTimestamp) { ])") .ValueOrDie()); - auto [orc_reader_holder, read_batch] = GenerateOrcReadBatch(src_array); + auto read_batch = GenerateOrcReadBatch(src_array); ASSERT_OK_AND_ASSIGN( std::shared_ptr target_array, OrcAdapter::AppendBatch(arrow_type, read_batch.get(), arrow::default_memory_pool())); diff --git a/src/paimon/format/orc/orc_file_batch_reader.cpp b/src/paimon/format/orc/orc_file_batch_reader.cpp index a201839b3..5be49f15b 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader.cpp @@ -60,29 +60,26 @@ void CollectAllColumnIds(const ::orc::Type* type, std::vector* column_ OrcFileBatchReader::OrcFileBatchReader(std::unique_ptr<::orc::ReaderMetrics>&& reader_metrics, std::unique_ptr&& reader, const std::map& options, - const std::shared_ptr& arrow_pool, - const std::shared_ptr<::orc::MemoryPool>& orc_pool) + const std::shared_ptr& read_memory) : options_(options), - arrow_pool_(arrow_pool), - orc_pool_(orc_pool), + read_memory_(read_memory), reader_metrics_(std::move(reader_metrics)), reader_(std::move(reader)), metrics_(std::make_shared()) {} Result> OrcFileBatchReader::Create( - std::unique_ptr<::orc::InputStream>&& input_stream, const std::shared_ptr& pool, + std::unique_ptr<::orc::InputStream>&& input_stream, + const std::shared_ptr& read_memory, const std::map& options, int32_t batch_size) { assert(input_stream); std::string file_name = input_stream->getName(); try { ::orc::ReaderOptions reader_options; - if (pool == nullptr) { - return Status::Invalid("memory pool is nullptr"); + if (read_memory == nullptr) { + return Status::Invalid("read memory is nullptr"); } uint64_t natural_read_size = input_stream->getNaturalReadSize(); - auto orc_pool = std::make_shared(pool); - std::shared_ptr arrow_pool = GetArrowPool(pool); - reader_options.setMemoryPool(*orc_pool); + reader_options.setMemoryPool(*read_memory->orc_pool); std::unique_ptr<::orc::ReaderMetrics> reader_metrics; PAIMON_ASSIGN_OR_RAISE( @@ -102,9 +99,9 @@ Result> OrcFileBatchReader::Create( PAIMON_ASSIGN_OR_RAISE( std::unique_ptr reader_wrapper, OrcReaderWrapper::Create(std::move(reader), file_name, batch_size, natural_read_size, - options, arrow_pool, orc_pool)); + options, read_memory->arrow_pool, read_memory->orc_pool)); auto orc_file_batch_reader = std::unique_ptr(new OrcFileBatchReader( - std::move(reader_metrics), std::move(reader_wrapper), options, arrow_pool, orc_pool)); + std::move(reader_metrics), std::move(reader_wrapper), options, read_memory)); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> file_schema, orc_file_batch_reader->GetFileSchema()); PAIMON_RETURN_NOT_OK(orc_file_batch_reader->SetReadSchema( @@ -196,6 +193,7 @@ Result OrcFileBatchReader::NextBatch() { if (BatchReader::IsEofBatch(batch)) { previous_batch_row_count_ = 0; } else { + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(batch.first.get(), read_memory_)); previous_batch_row_count_ = batch.first->length; } return batch; diff --git a/src/paimon/format/orc/orc_file_batch_reader.h b/src/paimon/format/orc/orc_file_batch_reader.h index b48f7cdb0..08f5ef27b 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.h +++ b/src/paimon/format/orc/orc_file_batch_reader.h @@ -30,6 +30,8 @@ #include "arrow/type.h" #include "orc/OrcFile.hh" #include "orc/Reader.hh" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/format/orc/orc_memory_pool.h" #include "paimon/format/orc/orc_reader_wrapper.h" #include "paimon/memory/memory_pool.h" #include "paimon/predicate/predicate.h" @@ -41,11 +43,20 @@ class InputStream; namespace paimon::orc { +struct OrcReadMemory { + explicit OrcReadMemory(const std::shared_ptr& pool) + : arrow_pool(GetSharedArrowPool(pool)), orc_pool(std::make_shared(pool)) {} + + std::shared_ptr arrow_pool; + std::shared_ptr<::orc::MemoryPool> orc_pool; +}; + class OrcFileBatchReader : public PrefetchFileBatchReader { public: ~OrcFileBatchReader() override = default; static Result> Create( - std::unique_ptr<::orc::InputStream>&& input_stream, const std::shared_ptr& pool, + std::unique_ptr<::orc::InputStream>&& input_stream, + const std::shared_ptr& read_memory, const std::map& options, int32_t batch_size); // For timestamp type, precision info is missing from file @@ -60,8 +71,7 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { return reader_->SetReadRanges(read_ranges); } - // Important: output ArrowArray is allocated on arrow_pool_ whose lifecycle holds in - // OrcFileBatchReader. Therefore, we need to hold BatchReader when using output ArrowArray. + // The output ArrowArray retains both Arrow and ORC pools and can outlive this reader. Result NextBatch() override; Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override { @@ -113,8 +123,7 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { OrcFileBatchReader(std::unique_ptr<::orc::ReaderMetrics>&& reader_metrics, std::unique_ptr&& reader, const std::map& options, - const std::shared_ptr& arrow_pool, - const std::shared_ptr<::orc::MemoryPool>& orc_pool); + const std::shared_ptr& read_memory); static Result<::orc::RowReaderOptions> CreateRowReaderOptions( const ::orc::Type* src_type, const ::orc::Type* target_type, @@ -127,8 +136,7 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { std::vector* target_column_ids); std::map options_; - std::shared_ptr arrow_pool_; - std::shared_ptr<::orc::MemoryPool> orc_pool_; + std::shared_ptr read_memory_; std::unique_ptr<::orc::ReaderMetrics> reader_metrics_; std::unique_ptr reader_; diff --git a/src/paimon/format/orc/orc_file_batch_reader_test.cpp b/src/paimon/format/orc/orc_file_batch_reader_test.cpp index c39ba5166..79466fabe 100644 --- a/src/paimon/format/orc/orc_file_batch_reader_test.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader_test.cpp @@ -64,6 +64,7 @@ class OrcFileBatchReaderTest : public ::testing::Test, public: void SetUp() override { pool_ = GetDefaultPool(); + read_memory_ = std::make_shared(pool_); batch_size_ = 10; arrow::FieldVector fields = { @@ -79,14 +80,12 @@ class OrcFileBatchReaderTest : public ::testing::Test, } void TearDown() override {} - std::pair, std::shared_ptr> - ReadBatchWithCustomizedData(const std::shared_ptr& src_array, - int32_t write_batch_size, int32_t write_stripe_size, - int32_t write_row_index_stride, const arrow::Schema* read_schema, - const std::shared_ptr& predicate, - const std::optional& selection_bitmap, - int32_t read_batch_size, double dict_key_size_threshold, - bool enable_lazy_decoding) const { + std::shared_ptr ReadBatchWithCustomizedData( + const std::shared_ptr& src_array, int32_t write_batch_size, + int32_t write_stripe_size, int32_t write_row_index_stride, const arrow::Schema* read_schema, + const std::shared_ptr& predicate, + const std::optional& selection_bitmap, int32_t read_batch_size, + double dict_key_size_threshold, bool enable_lazy_decoding) const { arrow::Schema src_schema(src_array->type()->fields()); EXPECT_OK_AND_ASSIGN(std::unique_ptr<::orc::Type> orc_type, OrcAdapter::GetOrcType(src_schema)); @@ -135,9 +134,9 @@ class OrcFileBatchReaderTest : public ::testing::Test, auto orc_batch_reader = PrepareOrcFileBatchReader(std::move(orc_input_stream), options, read_schema, predicate, selection_bitmap, read_batch_size); - EXPECT_OK_AND_ASSIGN( - auto result, paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get())); - return std::make_pair(std::move(orc_batch_reader), result); + EXPECT_OK_AND_ASSIGN(auto result, paimon::test::ReadResultCollector::CollectResult( + std::move(orc_batch_reader))); + return result; } std::unique_ptr PrepareOrcFileBatchReader( @@ -162,10 +161,12 @@ class OrcFileBatchReaderTest : public ::testing::Test, std::unique_ptr<::orc::InputStream>&& in_stream, const std::map& options, const arrow::Schema* read_schema, const std::shared_ptr& predicate, - const std::optional& selection_bitmap, int32_t batch_size) const { - EXPECT_OK_AND_ASSIGN( - auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), pool_, options, batch_size)); + const std::optional& selection_bitmap, int32_t batch_size, + const std::shared_ptr& read_memory = nullptr) const { + EXPECT_OK_AND_ASSIGN(auto orc_batch_reader, + OrcFileBatchReader::Create(std::move(in_stream), + read_memory ? read_memory : read_memory_, + options, batch_size)); EXPECT_TRUE(orc_batch_reader); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); @@ -197,6 +198,7 @@ class OrcFileBatchReaderTest : public ::testing::Test, private: std::shared_ptr pool_; + std::shared_ptr read_memory_; int32_t batch_size_; std::shared_ptr struct_array_; }; @@ -241,7 +243,7 @@ TEST_F(OrcFileBatchReaderTest, TestReadBinaryWrittenFromBinaryAndLargeBinary) { .ValueOrDie()); auto expected_chunked_array = std::make_shared(expected_array); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - orc_batch_reader.get())); + std::move(orc_batch_reader))); ASSERT_TRUE(result_array->Equals(expected_chunked_array)); }; @@ -260,7 +262,7 @@ TEST_F(OrcFileBatchReaderTest, TestSetReadSchema) { std::map options = {{ORC_READ_ENABLE_LAZY_DECODING, "true"}}; ASSERT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), pool_, options, batch_size_)); + OrcFileBatchReader::Create(std::move(in_stream), read_memory_, options, batch_size_)); // test GetFileSchema() ASSERT_OK_AND_ASSIGN(auto c_file_schema, orc_batch_reader->GetFileSchema()); auto arrow_file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); @@ -298,8 +300,8 @@ TEST_F(OrcFileBatchReaderTest, TestSetReadSchema) { ASSERT_TRUE(arrow::ExportSchema(read_schema, c_read_schema.get()).ok()); ASSERT_OK(orc_batch_reader->SetReadSchema(c_read_schema.get(), predicate, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(result_with_read_schema, - paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get())); + ASSERT_OK_AND_ASSIGN(result_with_read_schema, paimon::test::ReadResultCollector::CollectResult( + std::move(orc_batch_reader))); ASSERT_FALSE(result_with_read_schema); } @@ -615,6 +617,52 @@ TEST_P(OrcFileBatchReaderTest, TestNextBatchSimple) { } } +TEST_F(OrcFileBatchReaderTest, TestBatchRetainsOrcReadMemory) { + auto dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + std::shared_ptr file_system = dir->GetFileSystem(); + std::string file_name = dir->Str() + "/dictionary.orc"; + arrow::FieldVector fields = {arrow::field("f0", arrow::utf8())}; + std::shared_ptr expected = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ + ["abc"], ["abc"], ["abc"], ["de"], ["de"], ["de"] + ])") + .ValueOrDie()); + std::shared_ptr schema = arrow::schema(fields); + WriteArray(file_system, file_name, expected, schema, + {{ORC_DICTIONARY_KEY_SIZE_THRESHOLD, "0.9"}}); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system->Open(file_name)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr in_stream, + OrcInputStreamImpl::Create(input_stream, /*natural_read_size=*/16)); + + std::shared_ptr read_memory = std::make_shared(pool_); + std::weak_ptr weak_read_memory = read_memory; + std::map options = {{ORC_READ_ENABLE_LAZY_DECODING, "true"}}; + std::unique_ptr reader = PrepareOrcFileBatchReader( + std::move(in_stream), options, schema.get(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, /*batch_size=*/6, read_memory); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + + reader.reset(); + read_memory.reset(); + ASSERT_FALSE(weak_read_memory.expired()); + + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr array = std::move(import_result).ValueOrDie(); + ASSERT_FALSE(weak_read_memory.expired()); + ASSERT_TRUE(array->ValidateFull().ok()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr converted, + paimon::test::DictArrayConverter::ConvertDictArray(array, arrow::default_memory_pool())); + array.reset(); + ASSERT_TRUE(converted->Equals(expected)); + converted.reset(); + ASSERT_TRUE(weak_read_memory.expired()); +} + TEST_P(OrcFileBatchReaderTest, TestNextBatchWithTargetSchema) { std::string file_name = paimon::test::GetDataDir() + "/orc/append_09.db/append_09/f1=10/bucket-1/" @@ -628,8 +676,8 @@ TEST_P(OrcFileBatchReaderTest, TestNextBatchWithTargetSchema) { auto orc_batch_reader = PrepareOrcFileBatchReader(file_name, &read_schema, batch_size_, natural_read_size); - ASSERT_OK_AND_ASSIGN(auto result_array, - paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( + std::move(orc_batch_reader))); auto expected_array = std::make_shared( arrow::StructArray::Make( {struct_array_->GetFieldByName("f0"), struct_array_->GetFieldByName("f1"), @@ -656,7 +704,7 @@ TEST_F(OrcFileBatchReaderTest, TestNextBatchWithOutofOrderTargetSchema) { std::map options = {{ORC_READ_ENABLE_LAZY_DECODING, "true"}}; ASSERT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), pool_, options, batch_size_)); + OrcFileBatchReader::Create(std::move(in_stream), read_memory_, options, batch_size_)); std::unique_ptr c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(read_schema, c_schema.get()).ok()); ASSERT_NOK_WITH_MSG(orc_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, @@ -678,8 +726,8 @@ TEST_P(OrcFileBatchReaderTest, TestNextBatchWithNullValue) { auto [natural_read_size, _] = GetParam(); auto orc_batch_reader = PrepareOrcFileBatchReader(file_name, &read_schema, batch_size_, natural_read_size); - ASSERT_OK_AND_ASSIGN(auto result_array, - paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( + std::move(orc_batch_reader))); std::shared_ptr expected_array; std::string json = R"([ ["Paul", 20, 1, null] @@ -714,7 +762,7 @@ TEST_F(OrcFileBatchReaderTest, TestNextBatchWithDictionary) { auto read_schema = arrow::schema(read_fields); auto expected_chunk_array = std::make_shared(src_array); auto check_result = [&](double dict_key_size_threshold, bool enable_lazy_decoding) { - auto [orc_reader_holder, target_array] = ReadBatchWithCustomizedData( + auto target_array = ReadBatchWithCustomizedData( src_array, /*write_batch_size=*/src_array->length(), /*write_stripe_size=*/-1, /*write_row_index_stride=*/-1, read_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, 10, dict_key_size_threshold, enable_lazy_decoding); @@ -743,8 +791,8 @@ TEST_P(OrcFileBatchReaderTest, TestComplexType) { auto [natural_read_size, _] = GetParam(); auto orc_batch_reader = PrepareOrcFileBatchReader(file_name, &read_schema, batch_size_, natural_read_size); - ASSERT_OK_AND_ASSIGN(auto result_array, - paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( + std::move(orc_batch_reader))); std::shared_ptr expected_array; auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ [10, 1, 1234, "2033-05-18 03:33:20.0", "123456789987654321.45678", "add"], @@ -771,7 +819,7 @@ TEST_F(OrcFileBatchReaderTest, TestGetFileSchemaWithFieldId) { std::map options = {{ORC_READ_ENABLE_LAZY_DECODING, "true"}}; EXPECT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), pool_, options, batch_size_)); + OrcFileBatchReader::Create(std::move(in_stream), read_memory_, options, batch_size_)); EXPECT_TRUE(orc_batch_reader); auto c_file_schema = orc_batch_reader->GetFileSchema(); EXPECT_TRUE(c_file_schema.ok()); @@ -850,7 +898,7 @@ TEST_F(OrcFileBatchReaderTest, TestDictionaryWithMultiStripe) { .ValueOrDie()); auto src_schema = arrow::schema(fields); // force generate two stripe in one file - auto [orc_reader_holder, target_array] = ReadBatchWithCustomizedData( + auto target_array = ReadBatchWithCustomizedData( src_array, /*write_batch_size=*/3, /*write_stripe_size=*/3, /*write_row_index_stride=*/3, /*read_schema=*/src_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, /*read_batch_size=*/10, @@ -945,7 +993,7 @@ TEST_P(OrcFileBatchReaderTest, TestTimestampType) { batch_size_, natural_read_size); ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(orc_batch_reader))); ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString() << std::endl << expected_array->ToString(); } @@ -979,7 +1027,7 @@ TEST_P(OrcFileBatchReaderTest, TestTimestampType) { // check array ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(orc_batch_reader))); ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); } } @@ -1020,7 +1068,7 @@ TEST_F(OrcFileBatchReaderTest, TestNestedFieldProjection) { PrepareOrcFileBatchReader(data_path, &read_schema, /*batch_size=*/10, DEFAULT_NATURAL_READ_SIZE); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - orc_batch_reader.get())); + std::move(orc_batch_reader))); auto expected = std::dynamic_pointer_cast( arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({read_col1}), R"([ @@ -1044,7 +1092,7 @@ TEST_F(OrcFileBatchReaderTest, TestNestedFieldProjection) { PrepareOrcFileBatchReader(data_path, &read_schema, /*batch_size=*/10, DEFAULT_NATURAL_READ_SIZE); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - orc_batch_reader.get())); + std::move(orc_batch_reader))); auto expected = std::dynamic_pointer_cast( arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({read_col1, col2}), R"([ @@ -1068,7 +1116,7 @@ TEST_F(OrcFileBatchReaderTest, TestNestedFieldProjection) { PrepareOrcFileBatchReader(data_path, &read_schema, /*batch_size=*/10, DEFAULT_NATURAL_READ_SIZE); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - orc_batch_reader.get())); + std::move(orc_batch_reader))); auto expected = std::dynamic_pointer_cast( arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({read_col1}), R"([ @@ -1119,7 +1167,7 @@ TEST_F(OrcFileBatchReaderTest, TestDeepNestedFieldProjection) { PrepareOrcFileBatchReader(data_path, &read_schema, /*batch_size=*/10, DEFAULT_NATURAL_READ_SIZE); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - orc_batch_reader.get())); + std::move(orc_batch_reader))); auto expected = std::dynamic_pointer_cast( arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({read_a}), R"([ @@ -1166,7 +1214,7 @@ TEST_F(OrcFileBatchReaderTest, TestNestedFieldProjectionWithListAndMap) { PrepareOrcFileBatchReader(data_path, &read_schema, /*batch_size=*/10, DEFAULT_NATURAL_READ_SIZE); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - orc_batch_reader.get())); + std::move(orc_batch_reader))); auto expected = std::dynamic_pointer_cast( arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({read_col1}), R"([ @@ -1188,7 +1236,7 @@ TEST_F(OrcFileBatchReaderTest, TestNestedFieldProjectionWithListAndMap) { PrepareOrcFileBatchReader(data_path, &read_schema, /*batch_size=*/10, DEFAULT_NATURAL_READ_SIZE); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - orc_batch_reader.get())); + std::move(orc_batch_reader))); auto expected = std::dynamic_pointer_cast( arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({read_col1, col2}), R"([ @@ -1236,7 +1284,8 @@ TEST_F(OrcFileBatchReaderTest, TestListStructPartialProjection) { OrcInputStreamImpl::Create(input_stream, DEFAULT_NATURAL_READ_SIZE)); ASSERT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), pool_, /*options=*/{}, /*batch_size=*/10)); + OrcFileBatchReader::Create(std::move(in_stream), read_memory_, /*options=*/{}, + /*batch_size=*/10)); std::unique_ptr c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(read_schema, c_schema.get()).ok()); ASSERT_NOK_WITH_MSG(orc_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, @@ -1323,8 +1372,8 @@ TEST_F(OrcFileBatchReaderTest, TestAddMetadataPerFieldMetadata) { ASSERT_EQ("percent", unit_val); // Also verify data integrity — read it back and compare content. - ASSERT_OK_AND_ASSIGN(auto result_array, - paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( + std::move(orc_batch_reader))); ASSERT_EQ(result_array->num_chunks(), 1); ASSERT_TRUE(data->Equals(*result_array->chunk(0))) << result_array->ToString(); } diff --git a/src/paimon/format/orc/orc_reader_builder.h b/src/paimon/format/orc/orc_reader_builder.h index e7ef54022..f32055109 100644 --- a/src/paimon/format/orc/orc_reader_builder.h +++ b/src/paimon/format/orc/orc_reader_builder.h @@ -33,10 +33,14 @@ namespace paimon::orc { class OrcReaderBuilder : public ReaderBuilder { public: OrcReaderBuilder(const std::map& options, int32_t batch_size) - : batch_size_(batch_size), pool_(GetDefaultPool()), options_(options) {} + : batch_size_(batch_size), + pool_(GetDefaultPool()), + read_memory_(std::make_shared(pool_)), + options_(options) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; + read_memory_ = std::make_shared(pool); return this; } @@ -51,12 +55,14 @@ class OrcReaderBuilder : public ReaderBuilder { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr input_stream, OrcInputStreamImpl::Create(path, natural_read_size)); - return OrcFileBatchReader::Create(std::move(input_stream), pool_, options_, batch_size_); + return OrcFileBatchReader::Create(std::move(input_stream), read_memory_, options_, + batch_size_); } private: int32_t batch_size_ = -1; std::shared_ptr pool_; + std::shared_ptr read_memory_; std::map options_; }; } // namespace paimon::orc diff --git a/src/paimon/format/orc/predicate_pushdown_test.cpp b/src/paimon/format/orc/predicate_pushdown_test.cpp index 76d1e7408..100c05b79 100644 --- a/src/paimon/format/orc/predicate_pushdown_test.cpp +++ b/src/paimon/format/orc/predicate_pushdown_test.cpp @@ -69,6 +69,7 @@ class PredicatePushdownTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); + read_memory_ = std::make_shared(pool_); batch_size_ = 10; arrow::FieldVector fields = { @@ -114,14 +115,14 @@ class PredicatePushdownTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN(auto in_stream, OrcInputStreamImpl::Create(in, DEFAULT_NATURAL_READ_SIZE)); ASSERT_OK_AND_ASSIGN(auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), pool_, + OrcFileBatchReader::Create(std::move(in_stream), read_memory_, /*options=*/{}, batch_size_)); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); ASSERT_TRUE(arrow_status.ok()); ASSERT_OK(orc_batch_reader->SetReadSchema(c_schema.get(), predicate, /*selection_bitmap=*/std::nullopt)); - auto result = paimon::test::ReadResultCollector::CollectResult(orc_batch_reader.get()); + auto result = paimon::test::ReadResultCollector::CollectResult(std::move(orc_batch_reader)); if (result_ok) { ASSERT_TRUE(result.ok()); // check result @@ -140,6 +141,7 @@ class PredicatePushdownTest : public ::testing::Test { private: std::shared_ptr pool_; + std::shared_ptr read_memory_; int32_t batch_size_; std::shared_ptr struct_array_; std::shared_ptr fs_; @@ -420,7 +422,7 @@ TEST_F(PredicatePushdownTest, TestPredicatePushdownWithNullLiteral) { ASSERT_OK_AND_ASSIGN(auto in_stream, OrcInputStreamImpl::Create(in, DEFAULT_NATURAL_READ_SIZE)); ASSERT_OK_AND_ASSIGN(auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), pool_, + OrcFileBatchReader::Create(std::move(in_stream), read_memory_, /*options=*/{}, batch_size_)); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp b/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp index daacaceee..d17379348 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp @@ -213,8 +213,8 @@ class PageFilteredRowGroupReaderTest : public ::testing::Test { ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), predicate, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(*out, - paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN( + *out, paimon::test::ReadResultCollector::CollectResult(std::move(batch_reader))); } /// Read back a Parquet file with a predicate, a bitmap, and page index filter enabled. @@ -246,8 +246,8 @@ class PageFilteredRowGroupReaderTest : public ::testing::Test { // consuming data pages. tracking_in->ClearReadAtRanges(); } - ASSERT_OK_AND_ASSIGN(*out, - paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN( + *out, paimon::test::ReadResultCollector::CollectResult(std::move(batch_reader))); if (tracking_in) { *read_at_ranges = tracking_in->GetReadAtRanges(); } diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index 7605c4242..95f5cb75c 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -40,6 +40,7 @@ #include "arrow/util/thread_pool.h" #include "fmt/format.h" #include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/options_utils.h" @@ -637,6 +638,7 @@ Result ParquetFileBatchReader::NextBatch() { std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); read_rows_ += array->length(); read_batch_count_++; diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.h b/src/paimon/format/parquet/parquet_file_batch_reader.h index 2b1097fb4..870223c2f 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.h +++ b/src/paimon/format/parquet/parquet_file_batch_reader.h @@ -95,9 +95,7 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { return reader_->SeekToRow(row_number); } - // Important: output ArrowArray is allocated on arrow_pool_ whose lifecycle holds in - // ParquetFileBatchReader. Therefore, we need to hold BatchReader when using output - // ArrowArray. + // The output ArrowArray retains arrow_pool_ and can outlive this reader. Result NextBatch() override; Result>> GenReadRanges( diff --git a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp index 1409f24bf..86eab1458 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -366,7 +366,7 @@ TEST_F(ParquetFileBatchReaderTest, TestReadBinaryWrittenFromBinaryAndLargeBinary .ValueOrDie()); auto expected_chunked_array = std::make_shared(expected_array); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - parquet_batch_reader.get())); + std::move(parquet_batch_reader))); ASSERT_TRUE(result_array->Equals(expected_chunked_array)); }; @@ -382,7 +382,7 @@ TEST_F(ParquetFileBatchReaderTest, TestSimple) { file_name, schema_, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, batch_size_); ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); std::shared_ptr expected_array = std::make_shared(struct_array_); ASSERT_TRUE(result_array->Equals(*expected_array, @@ -441,7 +441,7 @@ TEST_F(ParquetFileBatchReaderTest, TestSetReadSchema) { ASSERT_OK(parquet_batch_reader->SetReadSchema(c_read_schema.get(), predicate, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(result_with_read_schema, paimon::test::ReadResultCollector::CollectResult( - parquet_batch_reader.get())); + std::move(parquet_batch_reader))); ASSERT_FALSE(result_with_read_schema); } @@ -462,7 +462,7 @@ TEST_F(ParquetFileBatchReaderTest, TestSetReadSchemaWithLegacyParquetMissingFiel /*selection_bitmap=*/std::nullopt, batch_size_); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - parquet_batch_reader.get())); + std::move(parquet_batch_reader))); std::shared_ptr expected_array; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -507,8 +507,7 @@ TEST_F(ParquetFileBatchReaderTest, TestNextBatchWithTargetSchema) { PrepareParquetFileBatchReader(file_name, read_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, batch_size_); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - parquet_batch_reader.get())); - parquet_batch_reader->Close(); + std::move(parquet_batch_reader))); auto expected_read_array = arrow::StructArray::Make({struct_array_->field(4), struct_array_->field(9), struct_array_->field(10), struct_array_->field(12)}, @@ -529,8 +528,7 @@ TEST_F(ParquetFileBatchReaderTest, TestNextBatchWithOutofOrderTargetSchema) { PrepareParquetFileBatchReader(file_name, read_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, batch_size_); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - parquet_batch_reader.get())); - parquet_batch_reader->Close(); + std::move(parquet_batch_reader))); auto expected_read_array = arrow::StructArray::Make({struct_array_->field(12), struct_array_->field(10), struct_array_->field(9), struct_array_->field(4)}, @@ -571,7 +569,7 @@ TEST_F(ParquetFileBatchReaderTest, TestNextBatchWithDictionary) { /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); }; check_result(true); @@ -608,7 +606,7 @@ TEST_F(ParquetFileBatchReaderTest, TestNestedStructChildProjectionRecall) { /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - parquet_batch_reader.get())); + std::move(parquet_batch_reader))); std::shared_ptr expected_array; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( @@ -651,7 +649,7 @@ TEST_F(ParquetFileBatchReaderTest, TestReadSchemaWithMapSelectedKeysMetadata) { /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - parquet_batch_reader.get())); + std::move(parquet_batch_reader))); auto expected_array = arrow::ChunkedArray::Make({write_array}).ValueOrDie(); ASSERT_TRUE(result_array->Equals(expected_array)) << "expected: " << expected_array->ToString() << "\nactual: " << result_array->ToString(); @@ -801,7 +799,7 @@ TEST_F(ParquetFileBatchReaderTest, TestNestedTimestampSecondReadFromMilliFile) { ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); auto expected_array = arrow::ChunkedArray::Make({write_array}).ValueOrDie(); ASSERT_TRUE(result_array->Equals(expected_array)) << "expected: " << expected_array->ToString() << "\nactual: " << result_array->ToString(); @@ -1029,7 +1027,7 @@ TEST_F(ParquetFileBatchReaderTest, TestBitmapRowGroupPushDownWithMultiRowGroups) ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); auto expected_array = arrow::ChunkedArray::Make({src_array->Slice(0, 6)}).ValueOrDie(); ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); @@ -1068,7 +1066,7 @@ TEST_F(ParquetFileBatchReaderTest, TestBitmapPagePushDownWithMultiRowGroups) { ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); auto expected_array = arrow::ChunkedArray(src_array->Slice(3, 3)); ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); @@ -1111,7 +1109,7 @@ TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapRowGroupPushDown) { file_path_, arrow_schema, predicate, bitmap, /*batch_size=*/length); ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); auto expected_array = arrow::ChunkedArray::Make({src_array->Slice(0, 256), src_array->Slice(512, 256)}) @@ -1127,7 +1125,7 @@ TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapRowGroupPushDown) { file_path_, arrow_schema, predicate, bitmap, /*batch_size=*/length); ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); ASSERT_FALSE(result_array); } } @@ -1169,7 +1167,7 @@ TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapPagePushDown) { /*batch_size=*/length, /*enable_page_level_filter=*/true); ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); auto expected_array = arrow::ChunkedArray::Make({src_array->Slice(100, 1), src_array->Slice(600, 1)}) @@ -1186,7 +1184,7 @@ TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapPagePushDown) { /*batch_size=*/length, /*enable_page_level_filter=*/true); ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); ASSERT_FALSE(result_array); } } @@ -1276,7 +1274,7 @@ TEST_P(ParquetFileBatchReaderTest, TestTimestampType) { /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, batch_size_); ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); } { @@ -1311,7 +1309,7 @@ TEST_P(ParquetFileBatchReaderTest, TestTimestampType) { // check array ASSERT_OK_AND_ASSIGN( std::shared_ptr result_array, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(parquet_batch_reader))); ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); } } @@ -1397,7 +1395,7 @@ TEST_F(ParquetFileBatchReaderTest, TestAddMetadataPerFieldMetadata) { // Also verify data integrity — read it back and compare content. ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( - parquet_batch_reader.get())); + std::move(parquet_batch_reader))); ASSERT_EQ(result_array->num_chunks(), 1); ASSERT_TRUE(data->Equals(*result_array->chunk(0))) << result_array->ToString(); } @@ -1709,7 +1707,7 @@ TEST_F(ParquetFileBatchReaderTest, TestPreBufferRangeFeedsReadAheadCache) { // Drain the file through the cache-backed stream. ASSERT_OK_AND_ASSIGN(std::shared_ptr result, - paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader)); + paimon::test::ReadResultCollector::CollectResult(std::move(base_reader))); ASSERT_EQ(20, result->length()); // The cache metrics must show that the data reads were served by the cache: diff --git a/src/paimon/format/parquet/parquet_reader_builder.h b/src/paimon/format/parquet/parquet_reader_builder.h index 112167bf3..394991369 100644 --- a/src/paimon/format/parquet/parquet_reader_builder.h +++ b/src/paimon/format/parquet/parquet_reader_builder.h @@ -49,10 +49,14 @@ namespace paimon::parquet { class ParquetReaderBuilder : public ReaderBuilder { public: ParquetReaderBuilder(const std::map& options, int32_t batch_size) - : batch_size_(batch_size), pool_(GetDefaultPool()), options_(options) {} + : batch_size_(batch_size), + pool_(GetDefaultPool()), + arrow_pool_(GetSharedArrowPool(pool_)), + options_(options) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; + arrow_pool_ = GetSharedArrowPool(pool); return this; } @@ -77,17 +81,16 @@ class ParquetReaderBuilder : public ReaderBuilder { file_uri = std::move(file_uri_result).value(); } } - std::shared_ptr arrow_pool = GetArrowPool(pool_); auto unique_input_stream = - std::make_unique(path, file_length, arrow_pool); + std::make_unique(path, file_length, arrow_pool_); auto storage_read_bytes = unique_input_stream->StorageReadBytes(); std::shared_ptr input_stream( std::move(unique_input_stream)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<::parquet::FileMetaData> file_metadata, - GetCachedParquetMetadata(input_stream, file_uri, arrow_pool)); + GetCachedParquetMetadata(input_stream, file_uri, arrow_pool_)); return ParquetFileBatchReader::Create( std::move(input_stream), options_, batch_size_, std::move(file_metadata), - std::move(storage_read_bytes), arrow_pool, hints_); + std::move(storage_read_bytes), arrow_pool_, hints_); } PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("ParquetReaderBuilder::Build") } @@ -167,6 +170,7 @@ class ParquetReaderBuilder : public ReaderBuilder { int32_t batch_size_ = -1; std::shared_ptr pool_; + std::shared_ptr arrow_pool_; std::map options_; std::shared_ptr cache_; std::optional hints_; diff --git a/src/paimon/format/parquet/parquet_vector_io_test.cpp b/src/paimon/format/parquet/parquet_vector_io_test.cpp index ad60caef3..db26f1b47 100644 --- a/src/paimon/format/parquet/parquet_vector_io_test.cpp +++ b/src/paimon/format/parquet/parquet_vector_io_test.cpp @@ -80,8 +80,9 @@ class ParquetVectorIoTest : public ::testing::Test { std::unique_ptr vector_reader; CreateVectorReader(file_path, arrow::schema(read_type->fields()), /*predicate=*/nullptr, /*options=*/{}, /*batch_size=*/10, &vector_reader); - ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - paimon::test::ReadResultCollector::CollectResult(vector_reader.get())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + paimon::test::ReadResultCollector::CollectResult(std::move(vector_reader))); arrow::Result> expected_result = arrow::ipc::internal::json::ArrayFromJSON(read_type, json); @@ -180,7 +181,7 @@ class ParquetVectorIoTest : public ::testing::Test { /*storage_read_bytes=*/nullptr, arrow_pool_, /*hints=*/std::nullopt)); std::unique_ptr vector_reader = - std::make_unique(std::move(reader), pool_); + std::make_unique(std::move(reader), arrow_pool_); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); ASSERT_OK(vector_reader->SetReadSchema(c_schema.get(), predicate, @@ -209,8 +210,9 @@ class ParquetVectorIoTest : public ::testing::Test { std::unique_ptr vector_reader; CreateVectorReader(file_path, logical_schema, /*predicate=*/nullptr, /*options=*/{}, /*batch_size=*/10, &vector_reader); - ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - paimon::test::ReadResultCollector::CollectResult(vector_reader.get())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + paimon::test::ReadResultCollector::CollectResult(std::move(vector_reader))); ASSERT_EQ(actual->num_chunks(), 1); ASSERT_EQ(actual->type()->id(), arrow::Type::STRUCT); auto struct_array = checked_pointer_cast(actual->chunk(0)); @@ -288,7 +290,7 @@ TEST_F(ParquetVectorIoTest, WriteAndReadAllNullFixedSizeListWithArrowSchema) { CreateVectorReader(file_path, arrow::schema(logical_type->fields()), /*predicate=*/nullptr, /*options=*/{}, /*batch_size=*/10, &reader); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); arrow::Result> expected_result = arrow::ipc::internal::json::ArrayFromJSON(logical_type, json); ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); @@ -350,7 +352,7 @@ TEST_F(ParquetVectorIoTest, ReadNestedFixedSizeListFile) { CreateVectorReader(file_path, arrow::schema(logical_type->fields()), /*predicate=*/nullptr, /*options=*/{}, /*batch_size=*/10, &reader); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); arrow::Result> expected_result = arrow::ipc::internal::json::ArrayFromJSON(logical_type, json); ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); @@ -377,7 +379,7 @@ TEST_F(ParquetVectorIoTest, ReadVectorWithPredicatePushdown) { CreateVectorReader(file_path, arrow::schema(logical_type->fields()), predicate, /*options=*/{}, /*batch_size=*/10, &reader); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - paimon::test::ReadResultCollector::CollectResult(reader.get())); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); arrow::Result> expected_result = arrow::ipc::internal::json::ArrayFromJSON( logical_type, R"([[3, [4.0, 5.0, 6.0]], [4, [7.0, 8.0, 9.0]]])"); @@ -427,10 +429,6 @@ TEST_F(ParquetVectorIoTest, ReadMixedListAndFixedSizeListFixtures) { arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3))}); std::shared_ptr logical_type = arrow::struct_(logical_schema->fields()); - // A reader owns the memory pool that its batches are allocated from, so it has to outlive - // the chunks collected from it. This mirrors a scan, which holds every split reader until - // the whole result has been consumed. - std::vector> readers; arrow::ArrayVector chunks; for (const char* file_name : {"java_vector_nullable.parquet", "rust_vector_nullable.parquet"}) { std::string file_path = @@ -439,8 +437,7 @@ TEST_F(ParquetVectorIoTest, ReadMixedListAndFixedSizeListFixtures) { CreateVectorReader(file_path, logical_schema, /*predicate=*/nullptr, /*options=*/{}, /*batch_size=*/10, &reader); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - paimon::test::ReadResultCollector::CollectResult(reader.get())); - readers.push_back(std::move(reader)); + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); ASSERT_TRUE(actual->type()->Equals(logical_type)) << file_name << ": " << actual->type()->ToString(); chunks.insert(chunks.end(), actual->chunks().begin(), actual->chunks().end()); diff --git a/src/paimon/format/parquet/predicate_pushdown_test.cpp b/src/paimon/format/parquet/predicate_pushdown_test.cpp index 26175e1df..03a0f69c1 100644 --- a/src/paimon/format/parquet/predicate_pushdown_test.cpp +++ b/src/paimon/format/parquet/predicate_pushdown_test.cpp @@ -137,8 +137,8 @@ class PredicatePushdownTest : public ::testing::Test { ASSERT_TRUE(arrow_status.ok()); ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), predicate, /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(auto arrow_array, - paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto arrow_array, paimon::test::ReadResultCollector::CollectResult( + std::move(batch_reader))); if (expected_array) { ASSERT_TRUE(arrow_array); auto expected_chunk_array = std::make_shared(expected_array); diff --git a/src/paimon/format/parquet/variant_parquet_test.cpp b/src/paimon/format/parquet/variant_parquet_test.cpp index 404b0521d..e0e73df15 100644 --- a/src/paimon/format/parquet/variant_parquet_test.cpp +++ b/src/paimon/format/parquet/variant_parquet_test.cpp @@ -423,8 +423,8 @@ class VariantParquetTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( read_schema, file_schema, pool_)); ASSERT_EQ(plans.size(), 1); - auto shredding_reader = - std::make_unique(std::move(file_reader), std::move(plans), pool_); + auto shredding_reader = std::make_unique( + std::move(file_reader), std::move(plans), arrow_pool_); auto c_read_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_read_schema.get()).ok()); ASSERT_OK(shredding_reader->SetReadSchema(c_read_schema.get(), /*predicate=*/nullptr, @@ -437,9 +437,7 @@ class VariantParquetTest : public ::testing::Test { auto result_struct = checked_pointer_cast(imported.ValueOrDie()); *column = result_struct->field(1); shredding_reader->Close(); - // The assembled arrays borrow the reader's memory pool; keep the reader alive until the - // fixture is torn down (fixture members outlive test-body locals). - live_readers_.push_back(std::move(shredding_reader)); + shredding_reader.reset(); } // `ReadColumn` for the cases whose second column is a struct. @@ -467,7 +465,6 @@ class VariantParquetTest : public ::testing::Test { std::shared_ptr map_item_field_; std::shared_ptr list_struct_schema_; std::shared_ptr list_struct_variant_field_; - std::vector> live_readers_; }; namespace { @@ -581,8 +578,7 @@ TEST_F(VariantParquetTest, WriteAndReadRoundTrip) { ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto result_chunked, - paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); - batch_reader->Close(); + paimon::test::ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_EQ(result_chunked->length(), static_cast(jsons.size())); ASSERT_EQ(result_chunked->num_chunks(), 1); auto result_struct = checked_pointer_cast(result_chunked->chunk(0)); diff --git a/src/paimon/testing/mock/mock_file_batch_reader.h b/src/paimon/testing/mock/mock_file_batch_reader.h index 4566289f4..866fc6369 100644 --- a/src/paimon/testing/mock/mock_file_batch_reader.h +++ b/src/paimon/testing/mock/mock_file_batch_reader.h @@ -116,8 +116,9 @@ class MockFileBatchReader : public PrefetchFileBatchReader { Result NextBatch() override { PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch_with_bitmap, NextBatchWithBitmap()); - return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), - arrow::default_memory_pool()); + std::shared_ptr arrow_pool(arrow::default_memory_pool(), + [](arrow::MemoryPool*) {}); + return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), arrow_pool); } Result NextBatchWithBitmap() override { diff --git a/src/paimon/testing/utils/read_result_collector.h b/src/paimon/testing/utils/read_result_collector.h index 558528662..219869777 100644 --- a/src/paimon/testing/utils/read_result_collector.h +++ b/src/paimon/testing/utils/read_result_collector.h @@ -31,6 +31,7 @@ #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/date_time_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/key_value_data_file_record_reader.h" #include "paimon/core/key_value.h" #include "paimon/reader/batch_reader.h" @@ -86,28 +87,48 @@ class ReadResultCollector { usleep(std::rand() % max_data_processing_time_in_us); } } - if (result_array_vector.empty()) { - return std::shared_ptr(); - } - // accumulate all the batch array and convert dictionary to string array together to avoid - // the problem (multiple batches in multiple stripes overlap dictionary data) being - // difficult to expose - arrow::ArrayVector converted_array_vector; - for (const auto& array : result_array_vector) { - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr converted_array, - DictArrayConverter::ConvertDictArray(array, arrow::default_memory_pool())); - converted_array_vector.push_back(converted_array); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(auto chunk_array, - arrow::ChunkedArray::Make(converted_array_vector)); - return chunk_array; + return BuildChunkedArray(result_array_vector); } static Result> CollectResult(BatchReader* batch_reader) { return CollectResult(batch_reader, /*max_data_processing_time_in_us=*/0); } + // Collect the C Arrow batches first, destroy the reader, and only then import and process + // them. Tests use this overload to verify that returned batches own every + // allocator resource needed by their release callbacks. + static Result> CollectResult( + std::unique_ptr batch_reader) { + std::vector batches; + // ReadBatch owns the C structs, but deleting those structs does not invoke their release + // callbacks. Explicitly release any buffered batches when a later read or import fails. + ScopeGuard release_batches([&batches]() { + for (BatchReader::ReadBatch& batch : batches) { + ReaderUtils::ReleaseReadBatch(std::move(batch)); + } + }); + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, + ReadOneRawBatch(batch_reader.get())); + if (BatchReader::IsEofBatch(batch)) { + break; + } + batches.push_back(std::move(batch)); + } + + batch_reader->Close(); + batch_reader.reset(); + + arrow::ArrayVector arrays; + arrays.reserve(batches.size()); + for (BatchReader::ReadBatch& batch : batches) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr array, + ImportReadBatch(std::move(batch))); + arrays.push_back(std::move(array)); + } + return BuildChunkedArray(arrays); + } + static Result> CollectResultOneBatch( BatchReader* batch_reader) { return CollectResultOneBatch(batch_reader, /*max_data_processing_time_in_us=*/0); @@ -175,6 +196,14 @@ class ReadResultCollector { private: static Result> ReadOneBatch(BatchReader* batch_reader) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, ReadOneRawBatch(batch_reader)); + if (BatchReader::IsEofBatch(batch)) { + return std::shared_ptr(); + } + return ImportReadBatch(std::move(batch)); + } + + static Result ReadOneRawBatch(BatchReader* batch_reader) { // Prioritize calling NextBatch. If it fails (paimon inner reader e.g., // PrefetchBatchReader, ApplyBitmapIndexBatchReader...), call NextBatchWithBitmap. auto batch_result = batch_reader->NextBatch(); @@ -185,25 +214,47 @@ class ReadResultCollector { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, batch_reader->NextBatchWithBitmap()); if (BatchReader::IsEofBatch(batch_with_bitmap)) { - return std::shared_ptr(); + return std::move(batch_with_bitmap.first); } PAIMON_RETURN_NOT_OK(CheckBatchOffset(batch_with_bitmap.first)); assert(!batch_with_bitmap.second.IsEmpty()); PAIMON_ASSIGN_OR_RAISE( - batch, ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), - arrow::default_memory_pool())); + batch, ReaderUtils::ApplyBitmapToReadBatch( + std::move(batch_with_bitmap), + std::shared_ptr(arrow::default_memory_pool(), + [](arrow::MemoryPool*) {}))); } else { return batch_result.status(); } } else { batch = std::move(batch_result).value(); if (BatchReader::IsEofBatch(batch)) { - return std::shared_ptr(); + return batch; } PAIMON_RETURN_NOT_OK(CheckBatchOffset(batch)); } assert(batch.first->length > 0); - return ImportReadBatch(std::move(batch)); + return batch; + } + + static Result> BuildChunkedArray( + const arrow::ArrayVector& arrays) { + if (arrays.empty()) { + return std::shared_ptr(); + } + // Accumulate all batches and convert dictionary arrays together to expose overlapping + // dictionaries from multiple stripes. + arrow::ArrayVector converted_arrays; + converted_arrays.reserve(arrays.size()); + for (const std::shared_ptr& array : arrays) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr converted_array, + DictArrayConverter::ConvertDictArray(array, arrow::default_memory_pool())); + converted_arrays.push_back(std::move(converted_array)); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr chunked_array, + arrow::ChunkedArray::Make(converted_arrays)); + return chunked_array; } static Result> ImportReadBatch(BatchReader::ReadBatch&& batch) { diff --git a/src/paimon/testing/utils/test_helper.h b/src/paimon/testing/utils/test_helper.h index 925a59ad3..10aa0a240 100644 --- a/src/paimon/testing/utils/test_helper.h +++ b/src/paimon/testing/utils/test_helper.h @@ -286,16 +286,17 @@ class TestHelper { PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr collected, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); if (collected->num_chunks() == 0) { return collected; } - // The collected batches borrow reader-owned buffers; copy them into the process pool - // while the reader is still alive so the returned result may outlive it. + // Preserve this helper's single-chunk result contract. CollectResult destroys the reader + // before importing the batches, so concatenating here also exercises the returned data + // after the reader has been released. PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr copied, + std::shared_ptr concatenated, arrow::Concatenate(collected->chunks(), arrow::default_memory_pool())); - return std::make_shared(copied); + return std::make_shared(std::move(concatenated)); } Result ReadAndCheckResult(const std::shared_ptr& data_type, @@ -308,7 +309,7 @@ class TestHelper { PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits)); PAIMON_ASSIGN_OR_RAISE(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( auto expected_array, arrow::ipc::internal::json::ArrayFromJSON(data_type, expected_result)); diff --git a/test/inte/append_compaction_inte_test.cpp b/test/inte/append_compaction_inte_test.cpp index ebeb23361..9c58f8b35 100644 --- a/test/inte/append_compaction_inte_test.cpp +++ b/test/inte/append_compaction_inte_test.cpp @@ -384,7 +384,8 @@ TEST_P(AppendCompactionInteTest, TestAppendTableStreamWriteFullCompactionWithMap ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto expected_type = arrow::struct_({ arrow::field("_VALUE_KIND", arrow::int8()), diff --git a/test/inte/blob_table_inte_test.cpp b/test/inte/blob_table_inte_test.cpp index f5d5960b3..30464f0b5 100644 --- a/test/inte/blob_table_inte_test.cpp +++ b/test/inte/blob_table_inte_test.cpp @@ -93,7 +93,6 @@ class RecordBatch; namespace paimon::test { struct ReadResult { - std::unique_ptr batch_reader; std::shared_ptr chunked_array; }; @@ -290,8 +289,7 @@ class BlobTableInteTest : public testing::Test, public ::testing::WithParamInter return result_plan; } - /// Read from table using a pre-scanned plan, returning the ChunkedArray and batch_reader. - /// The batch_reader must outlive the returned ChunkedArray (array memory depends on reader). + /// Read from table using a pre-scanned plan and return data after the reader is destroyed. Result ReadTable(const std::string& table_path, const std::vector& read_schema, const std::shared_ptr& plan, @@ -309,8 +307,8 @@ class BlobTableInteTest : public testing::Test, public ::testing::WithParamInter PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits)); PAIMON_ASSIGN_OR_RAISE(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); - return ReadResult{std::move(batch_reader), std::move(read_result)}; + ReadResultCollector::CollectResult(std::move(batch_reader))); + return ReadResult{std::move(read_result)}; } /// Convenience: scan + read in one call. diff --git a/test/inte/data_evolution_table_test.cpp b/test/inte/data_evolution_table_test.cpp index 303c6d700..4fb26b6ed 100644 --- a/test/inte/data_evolution_table_test.cpp +++ b/test/inte/data_evolution_table_test.cpp @@ -268,11 +268,6 @@ class DataEvolutionTableTest : public ::testing::Test, } struct LimitScanResult { - /// The read that produced `rows`. Its memory pool owns the buffers behind them, so it - /// has to outlive them: these two members are declared first on purpose, since members - /// are destroyed in reverse order. - std::unique_ptr table_read; - std::unique_ptr batch_reader; /// The splits the limit push down kept in the plan. std::vector> splits; /// The rows reading those splits produced, null when the read returned nothing. The @@ -312,17 +307,19 @@ class DataEvolutionTableTest : public ::testing::Test, .EnablePredicateFilter(enable_predicate_filter); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_context_builder.Finish()); - PAIMON_ASSIGN_OR_RAISE(result.table_read, TableRead::Create(std::move(read_context))); - PAIMON_ASSIGN_OR_RAISE(result.batch_reader, result.table_read->CreateReader(result.splits)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch_reader, + table_read->CreateReader(result.splits)); PAIMON_ASSIGN_OR_RAISE(result.rows, - ReadResultCollector::CollectResult(result.batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); return result; } /// Plans the table without any push down and returns the planned splits, so a test can /// assert what the scan handed the read: which data file each deletion file landed on, and /// the row count derived from them. Reading the splits back is ScanAndReadWithLimit's job, - /// which keeps the reader that owns the returned rows alive. + /// which verifies the returned rows after their reader has been destroyed. Result>> PlanSplits(const std::string& table_path) const { ScanContextBuilder scan_context_builder(table_path); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan_context, @@ -437,7 +434,7 @@ class DataEvolutionTableTest : public ::testing::Test, PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits)); PAIMON_ASSIGN_OR_RAISE(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); if (!expected_array) { if (read_result) { @@ -953,7 +950,8 @@ TEST_P(DataEvolutionTableTest, TestMultipleSharedShreddingMapsPartialOverwrite) read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto expected_type = arrow::struct_({ SpecialFields::ValueKind().field_, diff --git a/test/inte/global_index_test.cpp b/test/inte/global_index_test.cpp index 84ab22df8..b4ac8b83a 100644 --- a/test/inte/global_index_test.cpp +++ b/test/inte/global_index_test.cpp @@ -216,7 +216,7 @@ class GlobalIndexTest : public ::testing::Test, public ::testing::WithParamInter PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits)); PAIMON_ASSIGN_OR_RAISE(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); if (!expected_array) { if (read_result) { diff --git a/test/inte/nested_column_pruning_inte_test.cpp b/test/inte/nested_column_pruning_inte_test.cpp index b77d6bda9..119517af6 100644 --- a/test/inte/nested_column_pruning_inte_test.cpp +++ b/test/inte/nested_column_pruning_inte_test.cpp @@ -112,7 +112,8 @@ class NestedColumnPruningInteTest : public ::testing::Test, ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, + ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector expected_fields = expected_schema->fields(); expected_fields.insert(expected_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); @@ -780,7 +781,8 @@ TEST_P(NestedColumnPruningInteTest, PruneNestedStructWithSpecialFields) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_EQ(read_result->num_chunks(), 1); auto result_array = std::dynamic_pointer_cast(read_result->chunk(0)); @@ -1312,7 +1314,8 @@ TEST_P(NestedColumnPruningInteTest, MapSelectedKeysWithOrcDictionaryEncodedMap) ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_OK_AND_ASSIGN( auto decoded_result, diff --git a/test/inte/pk_compaction_inte_test.cpp b/test/inte/pk_compaction_inte_test.cpp index 8e9150aee..130535171 100644 --- a/test/inte/pk_compaction_inte_test.cpp +++ b/test/inte/pk_compaction_inte_test.cpp @@ -282,7 +282,7 @@ class PkCompactionInteTest : public ::testing::Test, ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, table_read->CreateReader(splits)); ASSERT_OK_AND_ASSIGN(std::shared_ptr read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto expected_array = arrow::ipc::internal::json::ArrayFromJSON(data_type, iter->second).ValueOrDie(); auto expected_chunk_array = std::make_shared(expected_array); @@ -299,9 +299,8 @@ class PkCompactionInteTest : public ::testing::Test, } } - // Read every row of the table, for fields whose expected value cannot be spelled out as JSON. - // `consume` runs while the reader is still alive, because the arrow arrays are allocated from - // a pool the reader owns and must not outlive it. + // Read every row of the table for fields whose expected value cannot be spelled out as JSON. + // The collector destroys the reader before `consume` accesses the returned arrays. template void ScanAllRows(const std::string& table_path, Fn consume) { std::map options = {{Options::FILE_SYSTEM, "local"}}; @@ -322,7 +321,7 @@ class PkCompactionInteTest : public ::testing::Test, ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, table_read->CreateReader(result_plan->Splits())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); consume(result); } @@ -1445,7 +1444,7 @@ TEST_F(PkCompactionInteTest, CompactWithSchemaEvolution) { { ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(split_p0)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1468,7 +1467,7 @@ TEST_F(PkCompactionInteTest, CompactWithSchemaEvolution) { { ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(split_p1)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -1967,7 +1966,7 @@ TEST_F(PkCompactionInteTest, WriteAndCompactWithBranch) { ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(std::shared_ptr(fake_split))); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), diff --git a/test/inte/primary_key_sorted_index_inte_test.cpp b/test/inte/primary_key_sorted_index_inte_test.cpp index 8c86f4e04..f158143ea 100644 --- a/test/inte/primary_key_sorted_index_inte_test.cpp +++ b/test/inte/primary_key_sorted_index_inte_test.cpp @@ -233,7 +233,7 @@ class PrimaryKeySortedIndexInteTest : public ::testing::Test, PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch_reader, table_read->CreateReader(splits)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); if (result == nullptr) { return std::vector(); } diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 343d99605..3b7a19654 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp @@ -220,13 +220,9 @@ class ReadInteTest : public testing::Test, public ::testing::WithParamInterface< namespace { struct SystemTableReadResult { - SystemTableReadResult(std::unique_ptr batch_reader, - std::shared_ptr array) - : batch_reader(std::move(batch_reader)), array(std::move(array)) {} + explicit SystemTableReadResult(std::shared_ptr array) + : array(std::move(array)) {} - // Keep BatchReader alive while the returned Arrow arrays are used. Some readers allocate - // exported ArrowArray buffers on pools owned by the BatchReader. - std::unique_ptr batch_reader; std::shared_ptr array; }; @@ -315,8 +311,8 @@ Result ReadSystemTable(const std::string& system_table_pa PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch_reader, table_read->CreateReader(plan->Splits())); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, - ReadResultCollector::CollectResult(batch_reader.get())); - return SystemTableReadResult(std::move(batch_reader), result); + ReadResultCollector::CollectResult(std::move(batch_reader))); + return SystemTableReadResult(std::move(result)); } Status WriteAndFullCompact(std::unique_ptr&& batch, int64_t commit_identifier, @@ -413,8 +409,24 @@ TEST_P(ReadInteTest, TestAppendSimple) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/3); ASSERT_EQ(data_splits.size(), 1); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); + + auto split_concat_batch_reader = dynamic_cast(batch_reader.get()); + ASSERT_TRUE(split_concat_batch_reader); + ASSERT_EQ(1, split_concat_batch_reader->readers_.size()); + auto complete_batch_reader = + dynamic_cast(split_concat_batch_reader->readers_[0].get()); + ASSERT_TRUE(complete_batch_reader); + auto file_concat_batch_reader = + dynamic_cast(complete_batch_reader->reader_.get()); + ASSERT_TRUE(file_concat_batch_reader); + ASSERT_EQ(2, file_concat_batch_reader->readers_.size()); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + std::shared_ptr read_metrics = batch_reader->GetReaderMetrics(); + std::shared_ptr reader0_metrics = + file_concat_batch_reader->readers_[0]->GetReaderMetrics(); + std::shared_ptr reader1_metrics = + file_concat_batch_reader->readers_[1]->GetReaderMetrics(); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -430,39 +442,17 @@ TEST_P(ReadInteTest, TestAppendSimple) { ASSERT_TRUE(array_status.ok()); ASSERT_TRUE(result_array->Equals(expected_array)); - // test metrics - auto read_metrics = batch_reader->GetReaderMetrics(); - - auto split_concat_batch_reader = dynamic_cast(batch_reader.get()); - ASSERT_TRUE(split_concat_batch_reader); - ASSERT_EQ(1, split_concat_batch_reader->readers_.size()); - auto complete_batch_reader = - dynamic_cast(split_concat_batch_reader->readers_[0].get()); - ASSERT_TRUE(complete_batch_reader); - auto file_concat_batch_reader = - dynamic_cast(complete_batch_reader->reader_.get()); - ASSERT_TRUE(file_concat_batch_reader); - ASSERT_EQ(2, file_concat_batch_reader->readers_.size()); - if (param.file_format == "orc") { - ASSERT_OK_AND_ASSIGN( - uint64_t reader0_latency, - file_concat_batch_reader->readers_[0]->GetReaderMetrics()->GetCounter( - "orc.read.inclusive.latency.us")); - ASSERT_OK_AND_ASSIGN( - uint64_t reader1_latency, - file_concat_batch_reader->readers_[1]->GetReaderMetrics()->GetCounter( - "orc.read.inclusive.latency.us")); + ASSERT_OK_AND_ASSIGN(uint64_t reader0_latency, + reader0_metrics->GetCounter("orc.read.inclusive.latency.us")); + ASSERT_OK_AND_ASSIGN(uint64_t reader1_latency, + reader1_metrics->GetCounter("orc.read.inclusive.latency.us")); uint64_t expected_read_latency = reader0_latency + reader1_latency; - ASSERT_OK_AND_ASSIGN( - uint64_t reader0_io_count, - file_concat_batch_reader->readers_[0]->GetReaderMetrics()->GetCounter( - "orc.read.io.count")); - ASSERT_OK_AND_ASSIGN( - uint64_t reader1_io_count, - file_concat_batch_reader->readers_[1]->GetReaderMetrics()->GetCounter( - "orc.read.io.count")); + ASSERT_OK_AND_ASSIGN(uint64_t reader0_io_count, + reader0_metrics->GetCounter("orc.read.io.count")); + ASSERT_OK_AND_ASSIGN(uint64_t reader1_io_count, + reader1_metrics->GetCounter("orc.read.io.count")); uint64_t expected_read_io_count = reader0_io_count + reader1_io_count; ASSERT_OK_AND_ASSIGN(uint64_t result_read_latency, @@ -523,18 +513,23 @@ TEST_P(ReadInteTest, TestReadWithLimits) { ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); // simulate read limits, only read 2 batches + std::vector batches; for (int32_t i = 0; i < 2; i++) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, batch_reader->NextBatch()); + batches.push_back(std::move(batch)); + } + batch_reader->Close(); + std::shared_ptr read_metrics = batch_reader->GetReaderMetrics(); + batch_reader.reset(); + + for (BatchReader::ReadBatch& batch : batches) { ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadResultCollector::GetArray(std::move(batch))); ASSERT_TRUE(array); ASSERT_EQ(array->length(), 1); } - batch_reader->Close(); - // test metrics if (param.file_format == "orc") { - auto read_metrics = batch_reader->GetReaderMetrics(); ASSERT_TRUE(read_metrics); ASSERT_OK_AND_ASSIGN(uint64_t io_count, read_metrics->GetCounter("orc.read.io.count")); ASSERT_GT(io_count, 0); @@ -576,13 +571,13 @@ TEST_P(ReadInteTest, TestReadAheadCacheMetrics) { ASSERT_EQ(data_splits.size(), 1); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + std::shared_ptr read_metrics = batch_reader->GetReaderMetrics(); ASSERT_TRUE(result_array); ASSERT_EQ(result_array->length(), 2); // Verify the read-ahead cache metrics are surfaced through the reader chain. The prefetch // reader merges the cache counters into its reader metrics only when a cache is created, // so the counters must be present and effective exactly in that case. - auto read_metrics = batch_reader->GetReaderMetrics(); ASSERT_TRUE(read_metrics); if (param.enable_prefetch && param.read_ahead_cache_enabled) { ASSERT_OK_AND_ASSIGN(uint64_t read_count, @@ -642,7 +637,8 @@ TEST_P(ReadInteTest, TestReadOnlyPartitionField) { } ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -693,7 +689,7 @@ TEST(SystemTableReadInteTest, TestReadOptionsSystemTable) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_TRUE(result); std::map expected = {{"custom.option", "custom-value"}, @@ -728,7 +724,7 @@ TEST(SystemTableReadInteTest, TestReadBranchOptionsSystemTable) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(std::move(batch_reader))); std::map expected = { {"bucket", "2"}, {"file.format", "parquet"}, {"manifest.format", "avro"}}; @@ -1058,7 +1054,7 @@ TEST(SystemTableReadInteTest, TestReadOptimizedPrimaryKeyProjectionAndPredicateP ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, table_read->CreateReader(ro_plan->Splits())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); std::shared_ptr expected_type = arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("k", arrow::int32()), arrow::field("v", arrow::int32())}); @@ -1119,7 +1115,7 @@ TEST(SystemTableReadInteTest, TestReadOptimizedSystemTableNestedProjection) { ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, table_read->CreateReader(plan->Splits())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); std::shared_ptr expected_type = arrow::struct_({ arrow::field("_VALUE_KIND", arrow::int8()), @@ -1990,7 +1986,8 @@ TEST_P(ReadInteTest, TestAppendReadWithMultipleBuckets) { BinaryRowGenerator::GenerateRow({20}, pool_.get()), file_list_2}}; auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/4); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2074,6 +2071,7 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicate) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/4); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + std::shared_ptr read_metrics = batch_reader->GetReaderMetrics(); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2088,10 +2086,8 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicate) { &expected_array); ASSERT_TRUE(array_status.ok()); ASSERT_TRUE(result_array->Equals(*expected_array)); - batch_reader->Close(); if (param.file_format == "orc") { // test metrics - auto read_metrics = batch_reader->GetReaderMetrics(); ASSERT_TRUE(read_metrics); ASSERT_OK_AND_ASSIGN(uint64_t io_count, read_metrics->GetCounter("orc.read.io.count")); ASSERT_GT(io_count, 0); @@ -2166,7 +2162,8 @@ TEST_P(ReadInteTest, TestAppendReadWithComplexTypePredicate) { BinaryRowGenerator::GenerateRow({20}, pool_.get()), file_list_2}}; auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/4); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2249,7 +2246,8 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicateOnlyPushdown) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/4); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2331,7 +2329,8 @@ TEST_P(ReadInteTest, TestAppendReadWithLateMaterializing) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/4); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2397,7 +2396,8 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicateAllFiltered) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/4); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_FALSE(result_array); } @@ -2462,7 +2462,8 @@ TEST_P(ReadInteTest, TestAppendReadIOException) { Result> batch_reader = table_read.value()->CreateReader(data_splits); CHECK_HOOK_STATUS(batch_reader.status(), i); - auto result = ReadResultCollector::CollectResult(batch_reader.value().get()); + std::unique_ptr owned_reader = std::move(batch_reader).value(); + auto result = ReadResultCollector::CollectResult(std::move(owned_reader)); CHECK_HOOK_STATUS(result.status(), i); auto result_array = result.value(); ASSERT_TRUE(result_array); @@ -2517,7 +2518,8 @@ TEST_P(ReadInteTest, TestPkTableWithDeletionVectorSimple) { std::shared_ptr arrow_data_type = DataField::ConvertDataFieldsToArrowStructType(fields_with_row_kind); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto expected = std::make_shared( arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([ [0, "Alex", 10, 0, 16.1], [0, "Bob", 10, 0, 12.1], @@ -2576,7 +2578,8 @@ TEST_P(ReadInteTest, TestPkTableWithDeletionVector) { /*snapshot_id=*/6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2663,7 +2666,8 @@ TEST_P(ReadInteTest, TestPkTableWithSnapshot6) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2736,7 +2740,8 @@ TEST_P(ReadInteTest, TestPkTableWithSnapshot8) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2807,7 +2812,7 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolution) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/2); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2903,7 +2908,8 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithPredicateFilter) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/2); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -2980,7 +2986,8 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithPredicateOnlyPushDown) auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/2); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3067,7 +3074,8 @@ TEST_P(ReadInteTest, TestPkReadSnapshot5WithSchemaEvolution) { // with new schema auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/5); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3143,7 +3151,8 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolution) { /*deletion file*/ {std::nullopt}}}; auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3229,7 +3238,8 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateOnlyPush /*deletion file*/ {std::nullopt}}}; auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3315,7 +3325,8 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithLateMaterializing auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3399,7 +3410,8 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateFilter) auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3467,7 +3479,8 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithBuildInFieldId) { /*schema ids*/ {0, 1}}}; auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/2); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3531,7 +3544,8 @@ TEST_P(ReadInteTest, TestAppendReadNestedType) { /*schema ids*/ {0}}}; auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/1); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3598,7 +3612,8 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithCast) { /*schema ids*/ {0, 1}}}; auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/2); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3682,7 +3697,8 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithCastWithPredicatePushD /*schema ids*/ {0, 1}}}; auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/2); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3738,7 +3754,7 @@ TEST_P(ReadInteTest, TestReadWithPKFallBackBranch) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3787,7 +3803,8 @@ TEST_P(ReadInteTest, TestReadWithAppendFallBackBranch) { ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); std::vector read_fields = { DataField(0, arrow::field("pt", arrow::int32())), @@ -3832,7 +3849,8 @@ TEST_P(ReadInteTest, TestFallBackBranchStreamRead) { ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_split)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3882,7 +3900,7 @@ TEST_P(ReadInteTest, TestReadWithPKRtBranch) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -3939,7 +3957,7 @@ TEST_P(ReadInteTest, TestReadWithAppendPtBranch) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -4084,7 +4102,8 @@ TEST_P(ReadInteTest, TestSpecificFs) { auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/3); ASSERT_EQ(data_splits.size(), 1); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -4141,8 +4160,8 @@ Result ReadGlobalSystemTable( PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch_reader, table_read->CreateReader(plan->Splits())); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, - ReadResultCollector::CollectResult(batch_reader.get())); - return SystemTableReadResult(std::move(batch_reader), result); + ReadResultCollector::CollectResult(std::move(batch_reader))); + return SystemTableReadResult(std::move(result)); } } // namespace diff --git a/test/inte/read_inte_with_index_test.cpp b/test/inte/read_inte_with_index_test.cpp index 1442f6263..515b812fc 100644 --- a/test/inte/read_inte_with_index_test.cpp +++ b/test/inte/read_inte_with_index_test.cpp @@ -94,7 +94,7 @@ class ReadInteWithIndexTest : public testing::Test, ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); ::arrow::PrettyPrintOptions print_option; print_option.container_window = 100; if (expected_array) { @@ -975,19 +975,18 @@ TEST_P(ReadInteWithIndexTest, TestSimple) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(split)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + std::shared_ptr read_metrics = batch_reader->GetReaderMetrics(); ASSERT_TRUE(result_array); ASSERT_TRUE(result_array->Equals(*expected_array)); // test metrics if (file_format == "orc") { - auto read_metrics = batch_reader->GetReaderMetrics(); ASSERT_OK_AND_ASSIGN(uint64_t io_count, read_metrics->GetCounter("orc.read.io.count")); ASSERT_GT(io_count, 0); ASSERT_OK_AND_ASSIGN(uint64_t latency, read_metrics->GetCounter("orc.read.inclusive.latency.us")); ASSERT_GT(latency, 0); } - batch_reader->Close(); } TEST_P(ReadInteWithIndexTest, TestReadWithLimits) { @@ -1043,18 +1042,24 @@ TEST_P(ReadInteWithIndexTest, TestReadWithLimits) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(split)); // simulate read limits, only read 3 batches + std::vector batches; for (int32_t i = 0; i < 3; i++) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, batch_reader->NextBatch()); + batches.push_back(std::move(batch)); + } + batch_reader->Close(); + std::shared_ptr read_metrics = batch_reader->GetReaderMetrics(); + batch_reader.reset(); + + for (BatchReader::ReadBatch& batch : batches) { ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadResultCollector::GetArray(std::move(batch))); ASSERT_TRUE(array); ASSERT_EQ(array->length(), 1); } - batch_reader->Close(); // test metrics if (file_format == "orc") { - auto read_metrics = batch_reader->GetReaderMetrics(); ASSERT_TRUE(read_metrics); ASSERT_OK_AND_ASSIGN(uint64_t io_count, read_metrics->GetCounter("orc.read.io.count")); ASSERT_GT(io_count, 0); @@ -1294,7 +1299,8 @@ TEST_P(ReadInteWithIndexTest, TestBitmapIndexWithLateMaterializing) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(std::vector>{split})); - ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto result_array, + ReadResultCollector::CollectResult(std::move(batch_reader))); // Only the two "Bob" rows match the predicate. std::shared_ptr expected_array; @@ -2601,7 +2607,8 @@ TEST_P(ReadInteWithIndexTest, TestWithIOException) { CHECK_HOOK_STATUS(table_read.status(), i); Result> batch_reader = table_read.value()->CreateReader(split); CHECK_HOOK_STATUS(batch_reader.status(), i); - auto result = ReadResultCollector::CollectResult(batch_reader.value().get()); + std::unique_ptr owned_reader = std::move(batch_reader).value(); + auto result = ReadResultCollector::CollectResult(std::move(owned_reader)); CHECK_HOOK_STATUS(result.status(), i); auto result_array = result.value(); ASSERT_TRUE(result_array); diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 4e28f286f..5362f664c 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -316,7 +316,6 @@ class RealtimeWriteInteTest : public ::testing::Test { using Row = std::tuple; struct CollectedReadResult { - std::unique_ptr reader; std::shared_ptr data; }; @@ -633,8 +632,8 @@ class RealtimeWriteInteTest : public ::testing::Test { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, table_read->CreateReader(plan->Splits())); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, - ReadResultCollector::CollectResult(reader.get())); - return CollectedReadResult{std::move(reader), std::move(result)}; + ReadResultCollector::CollectResult(std::move(reader))); + return CollectedReadResult{std::move(result)}; } void ReadPlanWithSchemaAndCheck(const std::shared_ptr& plan, @@ -654,7 +653,7 @@ class RealtimeWriteInteTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, table_read->CreateReader(plan->Splits())); ASSERT_OK_AND_ASSIGN(std::shared_ptr result, - ReadResultCollector::CollectResult(reader.get())); + ReadResultCollector::CollectResult(std::move(reader))); arrow::FieldVector result_fields = {arrow::field("_VALUE_KIND", arrow::int8())}; result_fields.insert(result_fields.end(), read_schema->fields().begin(), @@ -3027,7 +3026,7 @@ TEST_F(RealtimeWriteInteTest, TestReaderPinsMemoryAcrossRefresh) { ASSERT_LT(memory_usage_after_refresh, memory_usage_before_refresh); ASSERT_OK_AND_ASSIGN(std::shared_ptr result, - ReadResultCollector::CollectResult(reader.get())); + ReadResultCollector::CollectResult(std::move(reader))); std::shared_ptr result_type = arrow::struct_( {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::utf8())}); diff --git a/test/inte/scan_and_read_inte_test.cpp b/test/inte/scan_and_read_inte_test.cpp index 11080c9d5..de2d347ff 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -92,7 +92,7 @@ class ScanAndReadInteTest : public testing::Test, auto splits = result_plan->Splits(); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); if (expected_array[scan_id]) { ASSERT_TRUE(read_result); ASSERT_TRUE(expected_array[scan_id]->type()->Equals(read_result->type())); @@ -255,7 +255,8 @@ TEST_P(ScanAndReadInteTest, TestWithAppendSnapshotIOException) { Result> batch_reader = table_read.value()->CreateReader(splits); CHECK_HOOK_STATUS(batch_reader.status(), i); - auto read_result = ReadResultCollector::CollectResult(batch_reader.value().get()); + std::unique_ptr owned_reader = std::move(batch_reader).value(); + auto read_result = ReadResultCollector::CollectResult(std::move(owned_reader)); CHECK_HOOK_STATUS(read_result.status(), i); // check result @@ -310,7 +311,8 @@ TEST_P(ScanAndReadInteTest, TestWithPkSnapshotIOException) { Result> batch_reader = table_read.value()->CreateReader(splits); CHECK_HOOK_STATUS(batch_reader.status(), i); - auto read_result = ReadResultCollector::CollectResult(batch_reader.value().get()); + std::unique_ptr owned_reader = std::move(batch_reader).value(); + auto read_result = ReadResultCollector::CollectResult(std::move(owned_reader)); CHECK_HOOK_STATUS(read_result.status(), i); // check result @@ -355,7 +357,8 @@ TEST_P(ScanAndReadInteTest, TestWithAppendSnapshot1) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -396,7 +399,8 @@ TEST_P(ScanAndReadInteTest, TestWithAppendSnapshot3) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -438,7 +442,8 @@ TEST_P(ScanAndReadInteTest, TestWithAppendSnapshot5) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -534,7 +539,8 @@ TEST_P(ScanAndReadInteTest, TestJavaPaimon1WithAppendSnapshot1) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -589,7 +595,8 @@ TEST_P(ScanAndReadInteTest, TestJavaPaimon1WithAppendSnapshotOfNestedType) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result std::shared_ptr expected_array; @@ -622,7 +629,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -687,7 +695,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPartitionAndBu ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -731,7 +740,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPredicate) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -771,7 +781,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithLateMaterializ ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result: "Lucy" (f3 = 14.1) does not match f3 > 18 and is filtered out. auto expected = std::make_shared( @@ -823,7 +834,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithLimit) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 6); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1100,7 +1112,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithNestedType) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 2); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto struct_inner_type = arrow::struct_({arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int32()), @@ -1150,7 +1163,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanLatestSnapshot) { ASSERT_EQ(result_plan->SnapshotId().value(), 5); auto splits = result_plan->Splits(); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1197,7 +1211,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot2) { ASSERT_EQ(result_plan->SnapshotId().value(), 2); auto splits = result_plan->Splits(); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1240,7 +1255,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithPartitionAndB ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 5); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1297,7 +1313,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithPredicate) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 5); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1353,7 +1370,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithLateMateriali ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 5); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result: only the rows before "Lucy" with f3 <= 30.0 remain. auto expected = std::make_shared( @@ -1411,7 +1429,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvWithInvalidAggregateBatchScanSnapsho ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 3); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1481,7 +1500,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithAggregateBatchScanSnapshot3WithPredica ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 3); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1524,7 +1544,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithPartialUpdateBatchScanSnapshot3WithPre ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 3); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_TRUE(read_result); // check result auto expected = std::make_shared( @@ -1560,7 +1581,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithLimit) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 5); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1825,7 +1847,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithFirstRowBatchScanSnapshot5) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 5); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1956,7 +1979,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWith09VersionDvBatchScanLatestSnapshot) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_EQ(result_plan->SnapshotId().value(), 8); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -1998,7 +2022,7 @@ TEST_P(ScanAndReadInteTest, TestWithEmptyPartitionValue) { ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); }; @@ -2067,7 +2091,7 @@ TEST_P(ScanAndReadInteTest, TestWithMultipleEmptyPartitionValue) { ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); }; @@ -2121,7 +2145,7 @@ TEST_P(ScanAndReadInteTest, TestMemoryUse) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -2332,7 +2356,8 @@ TEST_P(ScanAndReadInteTest, TestScanWithPredicateAndReadWithUnorderedFieldForPar ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -2388,7 +2413,8 @@ TEST_P(ScanAndReadInteTest, TestPkSchemaEvolutionScanWithRenamedPkPredicate) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); auto expected = std::make_shared( arrow::ipc::internal::json::ArrayFromJSON( @@ -2432,7 +2458,8 @@ TEST_P(ScanAndReadInteTest, TestAppendTableWithMultipleFileFormat) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -2469,7 +2496,8 @@ TEST_P(ScanAndReadInteTest, TestPkDvTableIndexInDataAndNoExternalPath) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -2508,7 +2536,8 @@ TEST_P(ScanAndReadInteTest, TestPkDvTableIndexNotInDataAndNoExternalPath) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -2551,7 +2580,8 @@ TEST_P(ScanAndReadInteTest, TestPkDvTableIndexNotInDataAndWithExternalPath) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -2596,7 +2626,8 @@ TEST_P(ScanAndReadInteTest, TestScanAndReadWithDisableIndex) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result, when file-index.read.enabled = false, index will be ignored auto expected = std::make_shared( @@ -2639,7 +2670,8 @@ TEST_P(ScanAndReadInteTest, TestPkDvTableIndexInDataAndWithExternalPath) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto expected = std::make_shared( @@ -2678,7 +2710,8 @@ TEST_P(ScanAndReadInteTest, TestTimestampType) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto timezone = DateTimeUtils::GetLocalTimezoneName(); arrow::FieldVector fields = { @@ -2725,7 +2758,8 @@ TEST_P(ScanAndReadInteTest, TestCastTimestampType) { ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result arrow::FieldVector fields = { @@ -2822,7 +2856,7 @@ TEST_F(ScanAndReadInteTest, TestMosaicJavaAndPythonCompatibility) { ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, table_read->CreateReader(plan->Splits())); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_TRUE(expected_result->Equals(actual)) << "actual: " << (actual == nullptr ? "null" : actual->ToString()) << "\nexpected: " << expected_result->ToString(); @@ -2865,7 +2899,7 @@ TEST_F(ScanAndReadInteTest, TestAvroWithAppendTable) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result auto timezone = DateTimeUtils::GetLocalTimezoneName(); @@ -2939,7 +2973,7 @@ TEST_F(ScanAndReadInteTest, TestAvroWithPkTable) { ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); // check result arrow::FieldVector fields = { @@ -3018,7 +3052,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKBucketSelectByPredicate) { } ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(splits)); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // Only rows with f2=0 in partition f1=10 should be returned auto expected = std::make_shared( @@ -3061,11 +3096,19 @@ TEST_P(ScanAndReadInteTest, TestReadNullableMapKey) { const std::vector expected_rows = {R"([[0, 1, [["one", 10]]]])", R"([[0, 2, [["two", 20]]]])"}; + std::vector batches; for (int32_t i = 0; i < 2; ++i) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, batch_reader->NextBatch()); ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + batches.push_back(std::move(batch)); + } + ASSERT_NOK_WITH_MSG(batch_reader->NextBatch(), "Map array keys array should have no nulls"); + batch_reader->Close(); + batch_reader.reset(); + + for (int32_t i = 0; i < 2; ++i) { ASSERT_OK_AND_ASSIGN(std::shared_ptr array, - ReadResultCollector::GetArray(std::move(batch))); + ReadResultCollector::GetArray(std::move(batches[i]))); ASSERT_EQ(array->length(), 1); std::shared_ptr expected_array = @@ -3073,8 +3116,6 @@ TEST_P(ScanAndReadInteTest, TestReadNullableMapKey) { ASSERT_TRUE(array->Equals(expected_array)) << "actual: " << array->ToString() << ", expected: " << expected_array->ToString(); } - ASSERT_NOK_WITH_MSG(batch_reader->NextBatch(), "Map array keys array should have no nulls"); - batch_reader->Close(); } TEST_P(ScanAndReadInteTest, TestCountRowsEmptySplits) { @@ -3120,7 +3161,8 @@ TEST_P(ScanAndReadInteTest, TestCountRowsConsistencyWithCreateReader) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); int64_t iterate_count = read_result ? read_result->length() : 0; // Both methods should return the same count diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp index 43ed0d1fb..dae25488a 100644 --- a/test/inte/write_and_read_inte_test.cpp +++ b/test/inte/write_and_read_inte_test.cpp @@ -187,7 +187,8 @@ class WriteAndReadInteTest PAIMON_ASSIGN_OR_RAISE(auto read_context, read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(plan->Splits())); - PAIMON_ASSIGN_OR_RAISE(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + PAIMON_ASSIGN_OR_RAISE(auto actual, + ReadResultCollector::CollectResult(std::move(batch_reader))); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( auto expected, arrow::ipc::internal::json::ArrayFromJSON(expected_type, expected_data)); return std::make_shared(expected)->Equals(actual); @@ -209,7 +210,8 @@ class WriteAndReadInteTest PAIMON_ASSIGN_OR_RAISE(auto read_context, read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(plan->Splits())); - PAIMON_ASSIGN_OR_RAISE(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + PAIMON_ASSIGN_OR_RAISE(auto actual, + ReadResultCollector::CollectResult(std::move(batch_reader))); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( auto expected, arrow::ipc::internal::json::ArrayFromJSON(expected_type, expected_data)); return std::make_shared(expected)->Equals(actual); @@ -518,7 +520,7 @@ TEST_P(WriteAndReadInteTest, TestAppendVectorWithPredicate) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -603,7 +605,7 @@ TEST_P(WriteAndReadInteTest, TestAppendWithExternalBitmapAndRangeBitmapIndexes) ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -2319,7 +2321,8 @@ TEST_P(WriteAndReadInteTest, TestPKWithParquetPageIndexFilter) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // Expected: p2 file is pruned by file-level min/max key stats (f0 range // [Grace, Lucy] doesn't overlap "Alice"). Inside p1's file, write.batch-size=1 @@ -2428,7 +2431,8 @@ TEST_P(WriteAndReadInteTest, TestAppendWithParquetPageIndexFilter) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); // Partition p2's row groups don't overlap "Alice" (min/max f0 in [Grace, Lucy]), // so the whole file is skipped. Within p1, page-index pruning narrows down to the @@ -2521,7 +2525,8 @@ TEST_P(WriteAndReadInteTest, TestAppendWithParquetPageIndexFilterAndPrefetch) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -2602,7 +2607,7 @@ TEST_P(WriteAndReadInteTest, TestAppendWithParquetMetadataCache) { PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(result_plan->Splits())); PAIMON_ASSIGN_OR_RAISE(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); if (!read_result) { return Status::Invalid("read result is null"); } @@ -2952,7 +2957,7 @@ TEST_P(WriteAndReadInteTest, TestAppendMapSharedShreddingWithPredicate) { ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -4021,7 +4026,8 @@ TEST_P(WriteAndReadInteTest, TestMapSharedShreddingStructValueSchemaEvolutionRea PAIMON_ASSIGN_OR_RAISE(auto read_context, read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(plan->Splits())); - PAIMON_ASSIGN_OR_RAISE(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + PAIMON_ASSIGN_OR_RAISE(auto actual, + ReadResultCollector::CollectResult(std::move(batch_reader))); (void)actual; return Status::OK(); }; diff --git a/test/inte/write_inte_test.cpp b/test/inte/write_inte_test.cpp index 5b0b8a025..63768416c 100644 --- a/test/inte/write_inte_test.cpp +++ b/test/inte/write_inte_test.cpp @@ -295,7 +295,7 @@ class WriteInteTest : public testing::Test, public ::testing::WithParamInterface PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(data_splits)); PAIMON_ASSIGN_OR_RAISE(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); if (read_result == nullptr) { return Status::Invalid(fmt::format("No rows read for blob field {}", blob_field)); } @@ -2384,7 +2384,7 @@ TEST_F(WriteInteTest, TestPKTableWriteWithAlterTable) { ASSERT_OK(orc_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto result_array, - ReadResultCollector::CollectResult(orc_batch_reader.get())); + ReadResultCollector::CollectResult(std::move(orc_batch_reader))); std::shared_ptr expected_array; auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(read_fields), {R"([ @@ -2645,7 +2645,8 @@ TEST_P(WriteInteTest, TestAppendTableWriteAndReadWithExternalPath) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), @@ -2967,7 +2968,7 @@ TEST_P(WriteInteTest, TestWriteAndReadWithSpecialPartitionValue) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); // NOTE: // Users should not use the system-reserved keyword "__DEFAULT_PARTITION__" as a partition // value. If used, it may lead to behavioral inconsistencies between C++ Paimon and Java @@ -3017,7 +3018,7 @@ TEST_P(WriteInteTest, TestWriteAndReadWithSpecialPartitionValue) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); std::shared_ptr array = arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([ [0, "Alice", 10, 0, 11.1, " ", "a=b?"] @@ -3044,7 +3045,7 @@ TEST_P(WriteInteTest, TestWriteAndReadWithSpecialPartitionValue) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); std::shared_ptr array = arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([ [0, "Bob", 10, 0, 12.1, "", "a=b?"] @@ -3071,7 +3072,7 @@ TEST_P(WriteInteTest, TestWriteAndReadWithSpecialPartitionValue) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); ASSERT_OK_AND_ASSIGN(auto read_result, - ReadResultCollector::CollectResult(batch_reader.get())); + ReadResultCollector::CollectResult(std::move(batch_reader))); std::shared_ptr array = arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([ [0, "Cathy", 10, 0, 13.1, null, "a=b?"], @@ -3148,7 +3149,8 @@ TEST_P(WriteInteTest, TestWriteWithNestedSchema) { ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); - ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_OK_AND_ASSIGN(auto read_result, + ReadResultCollector::CollectResult(std::move(batch_reader))); arrow::FieldVector fields_with_row_kind = fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), From eb927ce83ca080d4bd97d0234ddd55eae8d63339 Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Sat, 29 Aug 2026 21:59:11 +0800 Subject: [PATCH 2/8] adjust tests --- ...hared_shredding_read_plan_factory_test.cpp | 8 +- .../apply_bitmap_index_batch_reader_test.cpp | 4 +- ...complete_index_score_batch_reader_test.cpp | 6 +- .../reader/blob_fallback_batch_reader.cpp | 9 +- .../reader/blob_fallback_batch_reader.h | 1 + .../blob_fallback_batch_reader_test.cpp | 39 ++- .../blob_view_resolving_batch_reader_test.cpp | 10 +- .../complete_row_kind_batch_reader_test.cpp | 6 +- .../common/reader/concat_batch_reader.cpp | 24 +- .../common/reader/concat_batch_reader.h | 3 + .../reader/concat_batch_reader_test.cpp | 58 +++- .../data_evolution_file_reader_test.cpp | 28 +- ...e_materializing_file_batch_reader_test.cpp | 69 +++-- .../prefetch_file_batch_reader_impl.cpp | 3 +- .../prefetch_file_batch_reader_impl_test.cpp | 290 +++++++++--------- ...mplete_row_tracking_fields_reader_test.cpp | 12 +- .../core/io/field_mapping_reader_test.cpp | 33 +- .../io/key_value_projection_reader_test.cpp | 10 +- .../compact/changelog_merge_tree_rewriter.cpp | 7 +- .../lookup_merge_tree_compact_rewriter.cpp | 4 +- .../compact/merge_tree_compact_rewriter.cpp | 18 +- .../compact/merge_tree_compact_rewriter.h | 4 +- src/paimon/core/mergetree/lookup_levels.cpp | 4 +- .../core/operation/abstract_split_read.cpp | 19 +- .../append_only_file_store_write.cpp | 7 +- .../core/operation/internal_read_context.cpp | 8 +- .../core/operation/internal_read_context.h | 3 +- .../operation/internal_read_context_test.cpp | 4 +- .../operation/merge_file_split_read_test.cpp | 10 +- .../operation/raw_file_split_read_test.cpp | 17 +- .../table/source/pk_count_reader_test.cpp | 4 +- src/paimon/core/table/source/table_read.cpp | 12 +- .../format/avro/avro_format_writer_test.cpp | 7 +- .../blob/blob_file_batch_reader_test.cpp | 83 ++--- .../format/blob/blob_format_writer_test.cpp | 73 ++--- src/paimon/format/orc/orc_reader_builder.h | 5 +- .../format/parquet/parquet_vector_io_test.cpp | 20 +- test/inte/blob_table_inte_test.cpp | 130 ++++---- test/inte/read_inte_test.cpp | 96 ++---- test/inte/realtime_write_inte_test.cpp | 69 ++--- 40 files changed, 610 insertions(+), 607 deletions(-) diff --git a/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp index 47e9649bd..a46986a12 100644 --- a/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp +++ b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp @@ -57,7 +57,6 @@ class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } static MapSharedShreddingFieldMeta TagsMeta() { @@ -139,7 +138,7 @@ class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test { field_read_plans.emplace(field->name(), std::move(field_read_plan)); } return std::make_unique(std::move(reader), std::move(field_read_plans), - arrow_pool_); + GetSharedArrowPool(pool_)); } Result> CreateReader( @@ -226,7 +225,6 @@ class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test { private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; std::shared_ptr logical_schema_ = arrow::schema({ arrow::field("id", arrow::int32()), arrow::field("tags", arrow::map(arrow::utf8(), arrow::int64())), @@ -305,7 +303,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjection) std::map> contexts; contexts.emplace("tags", std::move(field_read_plan)); auto reader = std::make_unique(std::move(mock_reader), std::move(contexts), - arrow_pool_); + GetSharedArrowPool(pool_)); auto read_schema = ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field})); @@ -445,7 +443,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjectionFr std::map> contexts; contexts.emplace("tags", std::move(field_read_plan)); auto reader = std::make_unique(std::move(mock_reader), std::move(contexts), - arrow_pool_); + GetSharedArrowPool(pool_)); auto read_schema = ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field})); diff --git a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp index cd713ab83..1b71c137a 100644 --- a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp +++ b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp @@ -56,7 +56,6 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, target_type_ = arrow::struct_({arrow::field("f1", int_type_)}); pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); fs_ = std::make_shared(); ASSERT_OK_AND_ASSIGN(executor_, CreateDefaultExecutor(/*thread_count=*/2)); } @@ -101,7 +100,7 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); @@ -125,7 +124,6 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, std::shared_ptr int_type_; std::shared_ptr target_type_; std::shared_ptr pool_; - std::shared_ptr arrow_pool_; std::shared_ptr fs_; std::shared_ptr executor_; }; diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp b/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp index 1d93fe2d6..6f1495c51 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp +++ b/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp @@ -41,7 +41,6 @@ class CompleteIndexScoreBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } void TearDown() override { pool_.reset(); @@ -53,7 +52,7 @@ class CompleteIndexScoreBatchReaderTest : public ::testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), selected_bitmap, batch_size); return std::make_unique(std::move(file_batch_reader), scores, - arrow_pool_); + GetSharedArrowPool(pool_)); } std::unique_ptr PrepareCompleteIndexScoreBatchReader( @@ -62,12 +61,11 @@ class CompleteIndexScoreBatchReaderTest : public ::testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), batch_size); return std::make_unique(std::move(file_batch_reader), scores, - arrow_pool_); + GetSharedArrowPool(pool_)); } private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_F(CompleteIndexScoreBatchReaderTest, TestSimple) { diff --git a/src/paimon/common/reader/blob_fallback_batch_reader.cpp b/src/paimon/common/reader/blob_fallback_batch_reader.cpp index bdd312238..a91b04852 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader.cpp +++ b/src/paimon/common/reader/blob_fallback_batch_reader.cpp @@ -100,7 +100,8 @@ BlobFallbackBatchReader::BlobFallbackBatchReader( row_id_field_idx_(row_id_field_idx), seq_num_field_idx_(seq_num_field_idx), read_batch_size_(read_batch_size), - arrow_pool_(arrow_pool) {} + arrow_pool_(arrow_pool), + finished_reader_metrics_(std::make_shared()) {} Result BlobFallbackBatchReader::FillWindow(size_t group_idx, int64_t want, std::vector* chunks) { @@ -156,6 +157,9 @@ Result BlobFallbackBatchReader::FillWindow(size_t group_idx, int64_t wa PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch_with_bitmap, segment.reader->NextBatchWithBitmap()); if (BatchReader::IsEofBatch(batch_with_bitmap)) { + segment.reader->Close(); + finished_reader_metrics_->Merge(segment.reader->GetReaderMetrics()); + segment.reader.reset(); cursor.segment_idx++; continue; } @@ -382,6 +386,8 @@ void BlobFallbackBatchReader::Close() { for (auto& segment : group.segments) { if (segment.reader) { segment.reader->Close(); + finished_reader_metrics_->Merge(segment.reader->GetReaderMetrics()); + segment.reader.reset(); } } } @@ -390,6 +396,7 @@ void BlobFallbackBatchReader::Close() { std::shared_ptr BlobFallbackBatchReader::GetReaderMetrics() const { auto metrics = std::make_shared(); + metrics->Merge(finished_reader_metrics_); for (const auto& group : groups_) { for (const auto& segment : group.segments) { if (segment.reader) { diff --git a/src/paimon/common/reader/blob_fallback_batch_reader.h b/src/paimon/common/reader/blob_fallback_batch_reader.h index 1e30b8957..011b3ff1b 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader.h +++ b/src/paimon/common/reader/blob_fallback_batch_reader.h @@ -158,6 +158,7 @@ class BlobFallbackBatchReader : public BatchReader { const int32_t seq_num_field_idx_; const int32_t read_batch_size_; std::shared_ptr arrow_pool_; + std::shared_ptr finished_reader_metrics_; bool closed_ = false; }; diff --git a/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp b/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp index 1750d8ef4..9b3fd5fe8 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp +++ b/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp @@ -46,7 +46,6 @@ class BlobFallbackBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); struct_type_ = arrow::struct_({BlobUtils::ToArrowField("blob_col", true)}); read_schema_ = arrow::schema(struct_type_->fields()); } @@ -119,9 +118,9 @@ class BlobFallbackBatchReaderTest : public ::testing::Test { for (const auto& specs : group_specs) { groups.push_back(MakeGroup(specs, file_batch_size)); } - ASSERT_OK_AND_ASSIGN( - auto reader, BlobFallbackBatchReader::Create(std::move(groups), read_schema_, - batch_size, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, BlobFallbackBatchReader::Create( + std::move(groups), read_schema_, batch_size, + GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result, paimon::test::ReadResultCollector::CollectResult( std::move(reader))); auto expected_chunk_array = std::make_shared(expected_array); @@ -135,7 +134,6 @@ class BlobFallbackBatchReaderTest : public ::testing::Test { protected: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; std::shared_ptr struct_type_; std::shared_ptr read_schema_; }; @@ -292,9 +290,9 @@ TEST_F(BlobFallbackBatchReaderTest, TestRowTrackingFieldsPreserved) { groups.push_back(std::move(newest)); groups.push_back(std::move(oldest)); - ASSERT_OK_AND_ASSIGN(auto reader, - BlobFallbackBatchReader::Create(std::move(groups), schema, - batch_size, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, BlobFallbackBatchReader::Create( + std::move(groups), schema, batch_size, + GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN( auto result, paimon::test::ReadResultCollector::CollectResult(std::move(reader))); @@ -320,8 +318,9 @@ TEST_F(BlobFallbackBatchReaderTest, TestMisalignedGroupsFail) { std::vector> groups; groups.push_back(MakeGroup({SegmentSpec::File({"PH", "u1", "PH"})}, 1024)); groups.push_back(MakeGroup({SegmentSpec::File({"b0", "b1"})}, 1024)); - ASSERT_OK_AND_ASSIGN(auto reader, BlobFallbackBatchReader::Create( - std::move(groups), read_schema_, 1024, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + BlobFallbackBatchReader::Create(std::move(groups), read_schema_, 1024, + GetSharedArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->NextBatch(), "same number of rows"); } @@ -329,9 +328,9 @@ TEST_F(BlobFallbackBatchReaderTest, TestCreateValidation) { // a single group needs no fallback std::vector> single_group; single_group.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); - ASSERT_NOK_WITH_MSG( - BlobFallbackBatchReader::Create(std::move(single_group), read_schema_, 1024, arrow_pool_), - "at least two sequence groups"); + ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(single_group), read_schema_, 1024, + GetSharedArrowPool(pool_)), + "at least two sequence groups"); // the read schema must contain a blob field std::vector> groups; @@ -339,25 +338,25 @@ TEST_F(BlobFallbackBatchReaderTest, TestCreateValidation) { groups.push_back(MakeGroup({SegmentSpec::File({"b1"})}, 1024)); auto plain_schema = arrow::schema({arrow::field("not_blob", arrow::large_binary(), /*nullable=*/true)}); - ASSERT_NOK_WITH_MSG( - BlobFallbackBatchReader::Create(std::move(groups), plain_schema, 1024, arrow_pool_), - "should contain a blob field"); + ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(groups), plain_schema, 1024, + GetSharedArrowPool(pool_)), + "should contain a blob field"); // groups must not be empty std::vector> with_empty_group; with_empty_group.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); with_empty_group.emplace_back(); ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(with_empty_group), read_schema_, - 1024, arrow_pool_), + 1024, GetSharedArrowPool(pool_)), "should not be empty"); // a gap segment must cover at least one selected row id std::vector> with_empty_gap; with_empty_gap.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); with_empty_gap.push_back(MakeGroup({SegmentSpec::GapRanges({})}, 1024)); - ASSERT_NOK_WITH_MSG( - BlobFallbackBatchReader::Create(std::move(with_empty_gap), read_schema_, 1024, arrow_pool_), - "at least one selected row id"); + ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(with_empty_gap), read_schema_, + 1024, GetSharedArrowPool(pool_)), + "at least one selected row id"); } } // namespace paimon::test diff --git a/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp b/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp index bdc4bbc6a..bd3f9c291 100644 --- a/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp +++ b/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp @@ -53,7 +53,6 @@ class BlobViewResolvingBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } void TearDown() override { @@ -132,7 +131,6 @@ class BlobViewResolvingBatchReaderTest : public ::testing::Test { private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_F(BlobViewResolvingBatchReaderTest, TestEofBatch) { @@ -141,7 +139,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestEofBatch) { return std::shared_ptr(); }); BlobViewResolvingBatchReader reader(std::move(inner_reader), {"blob_col"}, std::move(resolver), - arrow_pool_); + GetSharedArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto batch, reader.NextBatch()); ASSERT_TRUE(BatchReader::IsEofBatch(batch)); } @@ -160,7 +158,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestEmptyReadBlobViewFields) { auto inner_reader = std::make_unique(struct_array); auto reader = std::make_unique( std::move(inner_reader), /*read_blob_view_fields=*/std::vector(), - std::move(resolver), arrow_pool_); + std::move(resolver), GetSharedArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto expected_array = std::make_shared(struct_array); ASSERT_TRUE(expected_array->Equals(*result_array)); @@ -192,7 +190,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestResolvesBlobViewColumn) { auto inner_reader = std::make_unique(src_struct); auto reader = std::make_unique( std::move(inner_reader), std::vector{"blob_col"}, std::move(resolver), - arrow_pool_); + GetSharedArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto struct_array = std::dynamic_pointer_cast(result_array->chunk(0)); @@ -212,7 +210,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestResolverError) { }); auto inner_reader = std::make_unique(src_struct); BlobViewResolvingBatchReader reader(std::move(inner_reader), {"blob_col"}, std::move(resolver), - arrow_pool_); + GetSharedArrowPool(pool_)); ASSERT_NOK_WITH_MSG(reader.NextBatch(), "cache miss"); } diff --git a/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp b/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp index c12830102..7ae10be5c 100644 --- a/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp +++ b/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp @@ -47,7 +47,6 @@ class CompleteRowKindBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } void TearDown() override { @@ -70,7 +69,7 @@ class CompleteRowKindBatchReaderTest : public ::testing::Test { EXPECT_OK(orc_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); return std::make_unique(std::move(orc_batch_reader), - arrow_pool_); + GetSharedArrowPool(pool_)); } std::unique_ptr PrepareCompleteRowKindBatchReader( @@ -78,12 +77,11 @@ class CompleteRowKindBatchReaderTest : public ::testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), batch_size); return std::make_unique(std::move(file_batch_reader), - arrow_pool_); + GetSharedArrowPool(pool_)); } private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_F(CompleteRowKindBatchReaderTest, TestSimple) { diff --git a/src/paimon/common/reader/concat_batch_reader.cpp b/src/paimon/common/reader/concat_batch_reader.cpp index 9129d5702..5aea88e3a 100644 --- a/src/paimon/common/reader/concat_batch_reader.cpp +++ b/src/paimon/common/reader/concat_batch_reader.cpp @@ -31,7 +31,10 @@ class MemoryPool; ConcatBatchReader::ConcatBatchReader(std::vector>&& readers, const std::shared_ptr& arrow_pool) - : arrow_pool_(arrow_pool), readers_(std::move(readers)), current_(0) {} + : arrow_pool_(arrow_pool), + finished_reader_metrics_(std::make_shared()), + readers_(std::move(readers)), + current_(0) {} Result ConcatBatchReader::NextBatch() { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, @@ -44,12 +47,25 @@ Result ConcatBatchReader::NextBatch() { void ConcatBatchReader::Close() { for (; current_ < readers_.size(); current_++) { - readers_[current_]->Close(); + CloseAndReleaseReader(current_); } } std::shared_ptr ConcatBatchReader::GetReaderMetrics() const { - return MetricsImpl::CollectReadMetrics(readers_); + auto metrics = std::make_shared(); + metrics->Merge(finished_reader_metrics_); + metrics->Merge(MetricsImpl::CollectReadMetrics(readers_)); + return metrics; +} + +void ConcatBatchReader::CloseAndReleaseReader(size_t reader_index) { + std::unique_ptr& reader = readers_[reader_index]; + if (!reader) { + return; + } + reader->Close(); + finished_reader_metrics_->Merge(reader->GetReaderMetrics()); + reader.reset(); } Result ConcatBatchReader::NextBatchWithBitmap() { @@ -62,7 +78,7 @@ Result ConcatBatchReader::NextBatchWithBitmap( return result; } // current meets eof, move to next reader - current_reader->Close(); + CloseAndReleaseReader(current_); current_++; } // read finish diff --git a/src/paimon/common/reader/concat_batch_reader.h b/src/paimon/common/reader/concat_batch_reader.h index 4ce308aa3..2c889d44d 100644 --- a/src/paimon/common/reader/concat_batch_reader.h +++ b/src/paimon/common/reader/concat_batch_reader.h @@ -44,7 +44,10 @@ class ConcatBatchReader : public BatchReader { std::shared_ptr GetReaderMetrics() const override; private: + void CloseAndReleaseReader(size_t reader_index); + std::shared_ptr arrow_pool_; + std::shared_ptr finished_reader_metrics_; std::vector> readers_; size_t current_; }; diff --git a/src/paimon/common/reader/concat_batch_reader_test.cpp b/src/paimon/common/reader/concat_batch_reader_test.cpp index 85f716a75..a697e9ad6 100644 --- a/src/paimon/common/reader/concat_batch_reader_test.cpp +++ b/src/paimon/common/reader/concat_batch_reader_test.cpp @@ -29,6 +29,7 @@ #include "arrow/array/array_nested.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/memory/memory_pool.h" #include "paimon/status.h" @@ -38,10 +39,37 @@ #include "paimon/utils/roaring_bitmap32.h" namespace paimon::test { + +class LifetimeTrackingBatchReader : public BatchReader { + public: + LifetimeTrackingBatchReader(std::unique_ptr reader, + const std::shared_ptr& lifetime) + : reader_(std::move(reader)), lifetime_(lifetime) {} + + Result NextBatch() override { + return reader_->NextBatch(); + } + + Result NextBatchWithBitmap() override { + return reader_->NextBatchWithBitmap(); + } + + void Close() override { + reader_->Close(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + private: + std::unique_ptr reader_; + std::shared_ptr lifetime_; +}; + class ConcatBatchReaderTest : public ::testing::Test { void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } void CheckResult(const std::vector& batches, const std::string& expected) { std::vector>> batches_with_bitmap; @@ -74,7 +102,7 @@ class ConcatBatchReaderTest : public ::testing::Test { readers.push_back(std::move(reader)); } auto concat_reader = - std::make_unique(std::move(readers), arrow_pool_); + std::make_unique(std::move(readers), GetSharedArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_chunk_array, ReadResultCollector::CollectResult(std::move(concat_reader))); if (expected.empty()) { @@ -94,7 +122,6 @@ class ConcatBatchReaderTest : public ::testing::Test { private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_F(ConcatBatchReaderTest, TestSimple) { CheckResult({"[10, 11, 12, 13, 14]"}, "[10, 11, 12, 13, 14]"); @@ -173,4 +200,29 @@ TEST_F(ConcatBatchReaderTest, TestSimpleWithBitmap) { } } +TEST_F(ConcatBatchReaderTest, TestReleaseReaderAtEof) { + auto empty = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({arrow::field("f1", arrow::int32())}), "[]") + .ValueOrDie(); + auto data = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({arrow::field("f1", arrow::int32())}), "[[1]]") + .ValueOrDie(); + std::shared_ptr lifetime = std::make_shared(0); + std::weak_ptr weak_lifetime = lifetime; + + std::vector> readers; + readers.push_back(std::make_unique( + std::make_unique(empty, empty->type(), /*read_batch_size=*/1), + lifetime)); + readers.push_back( + std::make_unique(data, data->type(), /*read_batch_size=*/1)); + lifetime.reset(); + + ConcatBatchReader reader(std::move(readers), GetSharedArrowPool(pool_)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch, reader.NextBatchWithBitmap()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_TRUE(weak_lifetime.expired()); + ReaderUtils::ReleaseReadBatch(std::move(batch.first)); +} + } // namespace paimon::test diff --git a/src/paimon/common/reader/data_evolution_file_reader_test.cpp b/src/paimon/common/reader/data_evolution_file_reader_test.cpp index bae51a3f8..924ceadd3 100644 --- a/src/paimon/common/reader/data_evolution_file_reader_test.cpp +++ b/src/paimon/common/reader/data_evolution_file_reader_test.cpp @@ -40,7 +40,6 @@ class DataEvolutionFileReaderTest : public ::testing::Test, public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } void TearDown() override { @@ -75,10 +74,10 @@ class DataEvolutionFileReaderTest : public ::testing::Test, file_batch_reader->EnableRandomizeBatchSize(enable_randomize_batch_size); readers.push_back(std::move(file_batch_reader)); } - ASSERT_OK_AND_ASSIGN( - auto data_evolution_file_reader, - DataEvolutionFileReader::Create(std::move(readers), read_schema, batch_size, - reader_offsets, field_offsets, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto data_evolution_file_reader, + DataEvolutionFileReader::Create( + std::move(readers), read_schema, batch_size, reader_offsets, + field_offsets, GetSharedArrowPool(pool_))); // check metrics, data_evolution_file_reader collects all row of each // MockFileBatchReader auto metrics = data_evolution_file_reader->GetReaderMetrics(); @@ -112,7 +111,7 @@ class DataEvolutionFileReaderTest : public ::testing::Test, readers.push_back(std::move(file_batch_reader)); DataEvolutionFileReader fake_data_evolution_reader( std::move(readers), /*read_schema=*/arrow::schema({}), read_batch_size, - /*reader_offsets=*/{}, /*field_offsets=*/{}, arrow_pool_); + /*reader_offsets=*/{}, /*field_offsets=*/{}, GetSharedArrowPool(pool_)); arrow::ArrayVector result_array_vec; while (true) { ASSERT_OK_AND_ASSIGN(auto result_array, @@ -142,7 +141,6 @@ class DataEvolutionFileReaderTest : public ::testing::Test, private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_F(DataEvolutionFileReaderTest, TestInvalid) { @@ -150,7 +148,7 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { arrow::FieldVector read_fields; auto read_schema = arrow::schema(read_fields); ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, - {}, {}, arrow_pool_), + {}, {}, GetSharedArrowPool(pool_)), "read schema must not be empty"); } { @@ -165,7 +163,7 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { std::vector field_offsets = {0, 1, 0}; ASSERT_NOK_WITH_MSG( DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, reader_offsets, - field_offsets, arrow_pool_), + field_offsets, GetSharedArrowPool(pool_)), "read schema, row offsets and field offsets must have the same size"); } { @@ -180,7 +178,7 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { std::vector field_offsets = {0, 1, 1, 0}; ASSERT_NOK_WITH_MSG( DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, reader_offsets, - field_offsets, arrow_pool_), + field_offsets, GetSharedArrowPool(pool_)), "readers must not be empty"); } { @@ -196,10 +194,10 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { auto read_schema = arrow::schema(read_fields); std::vector reader_offsets = {0, 0, 1, 1}; std::vector field_offsets = {0, 1, 1, 0}; - ASSERT_NOK_WITH_MSG( - DataEvolutionFileReader::Create(std::move(readers), read_schema, /*read_batch_size=*/10, - reader_offsets, field_offsets, arrow_pool_), - "reader offset is out of range of readers"); + ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create( + std::move(readers), read_schema, /*read_batch_size=*/10, + reader_offsets, field_offsets, GetSharedArrowPool(pool_)), + "reader offset is out of range of readers"); } } @@ -640,7 +638,7 @@ TEST_P(DataEvolutionFileReaderTest, TestSingleReaderRowCountMismatch) { ASSERT_OK_AND_ASSIGN( auto data_evolution_file_reader, DataEvolutionFileReader::Create(std::move(readers), read_schema, /*read_batch_size=*/10, - reader_offsets, field_offsets, arrow_pool_)); + reader_offsets, field_offsets, GetSharedArrowPool(pool_))); // array0 has 6 rows but array1 only has 5 rows ASSERT_NOK_WITH_MSG( paimon::test::ReadResultCollector::CollectResult(std::move(data_evolution_file_reader)), diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index e75870157..f3aac3e93 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -60,7 +60,6 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); k_field_ = arrow::field("k", arrow::int64()); v_field_ = arrow::field("v", arrow::utf8()); full_fields_ = {k_field_, v_field_}; @@ -217,7 +216,6 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { protected: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; std::shared_ptr k_field_; std::shared_ptr v_field_; arrow::FieldVector full_fields_; @@ -228,8 +226,8 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr, std::nullopt)); @@ -247,8 +245,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); // read schema is just {k}; the predicate on k covers all columns -> payload empty auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(0l)); @@ -268,8 +266,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { TEST_F(LateMaterializingFileBatchReaderTest, ContiguousSubsetAcrossBatches) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(4l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -293,8 +291,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ScatteredAlternatingMatch) { } auto data = BuildData(ks); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -317,8 +315,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { } auto data = BuildData(ks); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); // predicate hits {1,3,5,7,9,11}; selection keeps only {1,5,9} @@ -343,8 +341,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(100l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -357,8 +355,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { TEST_F(LateMaterializingFileBatchReaderTest, SeekToRowRealignsProbeCursor) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(5l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -386,8 +384,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); auto* mock_ptr = mock.get(); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(2l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -415,8 +413,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); auto predicate1 = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(7l)); @@ -446,8 +444,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { TEST_F(LateMaterializingFileBatchReaderTest, ForwardsRowCountAndFileSchema) { auto data = BuildData({0, 1, 2, 3}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(uint64_t num_rows, reader->GetNumberOfRows()); EXPECT_EQ(num_rows, 4u); @@ -463,8 +461,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { auto data = BuildMultiFieldData(10); auto type = data->type(); auto mock = std::make_unique(data, type, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); // probe columns = {a (idx0), c (idx2)}; payload columns = {b, d, e} auto pred_a = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, Literal(3l)); @@ -507,8 +505,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { auto data = BuildNestedData(8); auto type = data->type(); auto mock = std::make_unique(data, type, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); // probe = {k}; payload = {arr (list), tag} auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); @@ -540,7 +538,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); LateMaterializingReaderBuilder builder( - std::make_unique(data, full_type_, /*batch_size=*/3), arrow_pool_); + std::make_unique(data, full_type_, /*batch_size=*/3), + GetSharedArrowPool(pool_)); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); ASSERT_OK_AND_ASSIGN( @@ -550,7 +549,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l)); ::ArrowSchema c_schema; @@ -575,7 +574,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); LateMaterializingReaderBuilder builder( - std::make_unique(data, full_type_, /*batch_size=*/3), arrow_pool_); + std::make_unique(data, full_type_, /*batch_size=*/3), + GetSharedArrowPool(pool_)); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); ASSERT_OK_AND_ASSIGN( @@ -585,7 +585,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto full_schema = arrow::schema(full_fields_); auto predicate1 = @@ -637,7 +637,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee // Per-batch ranges (the mock's default) let the impl split work across the parallel readers, // and each range-honoring reader only reads its assigned slice. LateMaterializingReaderBuilder builder( - std::make_unique(data, full_type_, /*batch_size=*/3), arrow_pool_); + std::make_unique(data, full_type_, /*batch_size=*/3), + GetSharedArrowPool(pool_)); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(3)); ASSERT_OK_AND_ASSIGN( @@ -647,7 +648,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee /*prefetch_max_parallel_num=*/3, /*batch_size=*/3, /*prefetch_batch_count=*/6, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); ::ArrowSchema c_schema; @@ -673,8 +674,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee TEST_F(LateMaterializingFileBatchReaderTest, FailsOnPredicateTypeMismatch) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, - LateMaterializingFileBatchReader::Create(std::move(mock), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( + std::move(mock), GetSharedArrowPool(pool_))); // k is int64 in the schema, but the predicate claims FieldType::INT (int32). auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::INT, Literal(10)); diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp index 859ceddd7..aaa2cd6e4 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -205,8 +205,7 @@ Result> PrefetchFileBatchReaderImpl uint32_t prefetch_batch_count, bool enable_adaptive_prefetch_strategy, const std::shared_ptr& executor, bool initialize_read_ranges, bool read_ahead_cache_enabled, const CacheConfig& cache_config, bool enable_io_metrics, - const std::shared_ptr& pool, - const std::shared_ptr& arrow_pool) { + const std::shared_ptr& pool, const std::shared_ptr& arrow_pool) { if (prefetch_max_parallel_num == 0) { return Status::Invalid("prefetch max parallel num should be greater than 0."); } diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp index 7e95e9f4d..d408b8641 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp @@ -181,7 +181,6 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); fields_ = {arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int64()), arrow::field("f2", arrow::boolean())}; data_type_ = arrow::struct_(fields_); @@ -270,7 +269,7 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, read_ahead_cache_enabled, CacheConfig(), - /*enable_io_metrics=*/true, pool_, arrow_pool_)); + /*enable_io_metrics=*/true, pool_, GetSharedArrowPool(pool_))); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); EXPECT_TRUE(arrow_status.ok()); @@ -306,7 +305,6 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; arrow::FieldVector fields_; std::shared_ptr data_type_; std::shared_ptr mock_fs_; @@ -356,12 +354,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestSimple) { for (auto prefetch_max_parallel_num : {1, 2, 3, 5, 8, 10}) { MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); ASSERT_OK_AND_ASSIGN( - auto reader, PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); if (prefetch_max_parallel_num == 1) { ASSERT_NOK( reader->GetReaderMetrics()->GetCounter(PrefetchIoMetrics::READ_LATENCY_COUNT)); @@ -382,13 +381,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLimits) { int32_t prefetch_max_parallel_num = 12; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/true, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/true, pool_, GetSharedArrowPool(pool_))); // simulate read limits, only read 8 batches for (int32_t i = 0; i < 8; i++) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, @@ -435,13 +435,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithoutInitializeReadRanges) { int32_t prefetch_max_parallel_num = 12; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); // simulate read limits, only read 8 batches ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "prefetch reader read ranges are not initialized"); @@ -459,7 +460,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestFailedIoMetrics) { /*prefetch_max_parallel_num=*/1, /*batch_size=*/10, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/true, pool_, arrow_pool_)); + /*enable_io_metrics=*/true, pool_, GetSharedArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "injected synchronous read failure"); std::shared_ptr metrics = reader->GetReaderMetrics(); @@ -541,13 +542,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRanges) { int32_t batch_size = 30; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_OK(prefetch_reader->RefreshReadRanges()); std::vector> read_ranges_0 = {{0, 30}, {90, 101}}; @@ -570,14 +572,15 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRangesDisablePrefetchByAdapti /*need_prefetch=*/true, /*set_read_ranges_statuses=*/{}); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, - /*prefetch_batch_count=*/2, - /*enable_adaptive_prefetch_strategy=*/true, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, + /*prefetch_batch_count=*/2, + /*enable_adaptive_prefetch_strategy=*/true, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); ASSERT_FALSE(reader->NeedPrefetch()); std::shared_ptr metrics = reader->GetReaderMetrics(); @@ -595,13 +598,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRanges) { int32_t batch_size = 30; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_FALSE(prefetch_reader->need_prefetch_); prefetch_reader->need_prefetch_ = true; @@ -638,13 +642,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRangesReturnErrorWhenPushDownFail /*set_read_ranges_statuses=*/ {Status::IOError("set read ranges failed"), Status::IOError("set read ranges failed")}); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->need_prefetch_ = true; @@ -665,12 +670,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, WorkloopSetReadStatusWhenCacheInitFailed /*pre_buffer_limit=*/128 * 1024); ASSERT_OK_AND_ASSIGN( - auto reader, PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - invalid_cache_config, /*enable_io_metrics=*/false, pool_, arrow_pool_)); + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + invalid_cache_config, /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->Workloop(); @@ -684,13 +690,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenShutdown) { int32_t batch_size = 5; int32_t prefetch_max_parallel_num = 1; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->is_shutdown_ = true; @@ -702,13 +709,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenNoCurrentReadRang int32_t batch_size = 5; int32_t prefetch_max_parallel_num = 1; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->read_ranges_in_group_ = {{}}; @@ -720,13 +728,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLargeBatchSize) { int32_t batch_size = 150; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -740,13 +749,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPartialReaderSuccessRead) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { dynamic_cast(prefetch_reader->readers_[i].get()) @@ -785,13 +795,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { @@ -823,13 +834,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithEmptyData) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -842,13 +854,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCallNextBatchAfterReadingEof) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 6; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -885,13 +898,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCreateReaderWithoutNextBatch) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); } TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { @@ -906,7 +920,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*prefetch_max_parallel_num=*/0, batch_size, 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -914,7 +928,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, /*batch_size=*/-1, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -923,7 +937,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*enable_adaptive_prefetch_strategy=*/false, /*executor=*/nullptr, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), /*enable_io_metrics=*/false, pool_, - arrow_pool_)); + GetSharedArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -931,7 +945,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -939,16 +953,17 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*fs=*/nullptr, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); } { ASSERT_OK_AND_ASSIGN( - auto reader, PrefetchFileBatchReaderImpl::Create( - data_file_path, /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, - CacheConfig(), /*enable_io_metrics=*/false, pool_, arrow_pool_)); + auto reader, + PrefetchFileBatchReaderImpl::Create( + data_file_path, /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SeekToRow(/*row_number=*/101), "not support seek to row for prefetch reader"); } @@ -1045,21 +1060,22 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithBitmap) { MockFormatReaderBuilder reader_builder(data_array, data_type_, bitmap, /*read_batch_size=*/100); int32_t prefetch_max_parallel_num = 3; - ASSERT_OK_AND_ASSIGN(auto reader, PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, - &reader_builder, mock_fs_, prefetch_max_parallel_num, - /*batch_size=*/100, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, - /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, + /*batch_size=*/100, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, + /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_chunk_array, ReadResultCollector::CollectResult(std::move(reader))); ASSERT_OK_AND_ASSIGN(auto data_batch, ReadResultCollector::GetReadBatch(data_array)); - ASSERT_OK_AND_ASSIGN(auto expected_batch, - ReaderUtils::ApplyBitmapToReadBatch( - std::make_pair(std::move(data_batch), bitmap), arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto expected_batch, ReaderUtils::ApplyBitmapToReadBatch( + std::make_pair(std::move(data_batch), bitmap), + GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto expected_array, ReadResultCollector::GetArray(std::move(expected_batch))); auto expected_chunk_array = std::make_shared(expected_array); diff --git a/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp b/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp index cad0a6c98..d070d2cbb 100644 --- a/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp +++ b/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp @@ -40,7 +40,6 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } void CheckResult(const std::shared_ptr& src_array, int64_t first_row_id, @@ -52,7 +51,7 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { auto complete_row_tracking_fields_reader = std::make_unique( std::move(file_batch_reader), first_row_id, snapshot_id, - /*file_field_names=*/std::nullopt, arrow_pool_); + /*file_field_names=*/std::nullopt, GetSharedArrowPool(pool_)); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK(complete_row_tracking_fields_reader->SetReadSchema( @@ -77,7 +76,7 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/10, /*snapshot_id=*/1, - file_field_names, arrow_pool_); + file_field_names, GetSharedArrowPool(pool_)); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( @@ -93,7 +92,6 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestSetReadSchema) { @@ -380,7 +378,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidWithReadNonExistFiel auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/100, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, arrow_pool_); + /*file_field_names=*/std::nullopt, GetSharedArrowPool(pool_)); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( @@ -417,7 +415,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidNextBatchBeforeSetRe auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/100, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, arrow_pool_); + /*file_field_names=*/std::nullopt, GetSharedArrowPool(pool_)); ASSERT_NOK_WITH_MSG(complete_row_tracking_fields_reader->NextBatchWithBitmap(), "in CompleteRowTrackingFieldsBatchReader SetReadSchema is supposed to be " "called before NextBatch"); @@ -448,7 +446,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidNullFirstRowId) { auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/std::nullopt, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, arrow_pool_); + /*file_field_names=*/std::nullopt, GetSharedArrowPool(pool_)); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( diff --git a/src/paimon/core/io/field_mapping_reader_test.cpp b/src/paimon/core/io/field_mapping_reader_test.cpp index 9a5a80bfa..19ac1c537 100644 --- a/src/paimon/core/io/field_mapping_reader_test.cpp +++ b/src/paimon/core/io/field_mapping_reader_test.cpp @@ -61,7 +61,6 @@ namespace paimon::test { class FieldMappingReaderTest : public ::testing::Test { public: void SetUp() override { - arrow_pool_ = GetSharedArrowPool(pool_); partition_ = BinaryRowGenerator::GenerateRow({10, 0}, pool_.get()); } void TearDown() override {} @@ -141,7 +140,8 @@ class FieldMappingReaderTest : public ::testing::Test { auto reader, FieldMappingReader::Create( /*field_count=*/read_schema->num_fields(), std::move(orc_batch_reader), partition_, - std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_)); + std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, + GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); if (expect_array == nullptr && result_array == nullptr) { @@ -183,7 +183,8 @@ class FieldMappingReaderTest : public ::testing::Test { auto reader, FieldMappingReader::Create( /*field_count=*/read_schema->num_fields(), std::move(orc_batch_reader), partition, - std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_)); + std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, + GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); if (expect_array == nullptr && result_array == nullptr) { @@ -202,7 +203,6 @@ class FieldMappingReaderTest : public ::testing::Test { private: std::shared_ptr pool_ = GetDefaultPool(); - std::shared_ptr arrow_pool_; std::vector data_fields_ = {DataField(0, arrow::field("f0", arrow::utf8())), DataField(1, arrow::field("f1", arrow::int32())), DataField(2, arrow::field("f2", arrow::int32())), @@ -247,7 +247,7 @@ TEST_F(FieldMappingReaderTest, TestGenerateSinglePartitionArray) { auto mapping_reader, FieldMappingReader::Create( /*field_count=*/8, /*reader=*/nullptr, partition, std::move(field_mapping), - /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_)); + /*skip_map_selected_keys_filter_field_ids=*/{}, GetSharedArrowPool(pool_))); { ASSERT_OK_AND_ASSIGN(auto p7_array, mapping_reader->GenerateSinglePartitionArray( @@ -475,11 +475,11 @@ TEST_F(FieldMappingReaderTest, TestSchemaEvolutionAddedFieldInsideList) { ASSERT_OK_AND_ASSIGN(auto mapping, mapping_builder->CreateFieldMapping(data_fields)); auto mock = std::make_unique( data_array, arrow::struct_(data_schema->fields()), /*read_batch_size=*/8); - ASSERT_OK_AND_ASSIGN( - auto reader, - FieldMappingReader::Create(read_schema->num_fields(), std::move(mock), - BinaryRow::EmptyRow(), std::move(mapping), - /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + FieldMappingReader::Create(read_schema->num_fields(), std::move(mock), + BinaryRow::EmptyRow(), std::move(mapping), + /*skip_map_selected_keys_filter_field_ids=*/{}, + GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto expect_array = @@ -865,11 +865,12 @@ TEST_F(FieldMappingReaderTest, TestCreateFailFastOnInvalidMapSelectedKeysMetadat /*partition_info=*/PartitionInfo(), std::move(non_part_info), /*non_exist_field_info=*/std::nullopt); - ASSERT_NOK_WITH_MSG(FieldMappingReader::Create( - /*field_count=*/1, - /*reader=*/nullptr, - /*partition=*/BinaryRow::EmptyRow(), std::move(field_mapping), - /*skip_map_selected_keys_filter_field_ids=*/{}, arrow_pool_), - "Duplicate selected key 'a'"); + ASSERT_NOK_WITH_MSG( + FieldMappingReader::Create( + /*field_count=*/1, + /*reader=*/nullptr, + /*partition=*/BinaryRow::EmptyRow(), std::move(field_mapping), + /*skip_map_selected_keys_filter_field_ids=*/{}, GetSharedArrowPool(pool_)), + "Duplicate selected key 'a'"); } } // namespace paimon::test diff --git a/src/paimon/core/io/key_value_projection_reader_test.cpp b/src/paimon/core/io/key_value_projection_reader_test.cpp index 6121698c1..1b15a77e1 100644 --- a/src/paimon/core/io/key_value_projection_reader_test.cpp +++ b/src/paimon/core/io/key_value_projection_reader_test.cpp @@ -58,7 +58,6 @@ class KeyValueProjectionReaderTest : public testing::Test, public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } std::unique_ptr GenerateProjectionReader( @@ -97,14 +96,14 @@ class KeyValueProjectionReaderTest : public testing::Test, if (!multi_thread_row_to_batch) { EXPECT_OK_AND_ASSIGN( auto projection_reader, - KeyValueProjectionReader::Create(std::move(sort_merge_reader), target_schema, - target_to_src_mapping, - /*batch_size=*/batch_size, arrow_pool_)); + KeyValueProjectionReader::Create( + std::move(sort_merge_reader), target_schema, target_to_src_mapping, + /*batch_size=*/batch_size, GetSharedArrowPool(pool_))); return std::move(projection_reader); } else { return std::make_unique( std::move(sort_merge_reader), target_schema, target_to_src_mapping, batch_size, - /*projection_thread_num=*/3, arrow_pool_); + /*projection_thread_num=*/3, GetSharedArrowPool(pool_)); } } @@ -159,7 +158,6 @@ class KeyValueProjectionReaderTest : public testing::Test, private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_P(KeyValueProjectionReaderTest, TestBulkData) { diff --git a/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.cpp b/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.cpp index 4f632425b..159a82cb6 100644 --- a/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.cpp @@ -227,7 +227,6 @@ Result ChangelogMergeTreeRewriter::RewriteOrProduceChangelog( PAIMON_ASSIGN_OR_RAISE(changelog_file_writer, CreateRollingChangelogWriter(output_level)); } - std::vector> reader_holders; ScopeGuard write_guard([&]() -> void { if (compact_file_writer) { compact_file_writer->Abort(); @@ -237,9 +236,6 @@ Result ChangelogMergeTreeRewriter::RewriteOrProduceChangelog( changelog_file_writer->Abort(); changelog_file_writer.reset(); } - for (const auto& reader : reader_holders) { - reader->Close(); - } merge_file_split_read_.reset(); }); @@ -263,9 +259,8 @@ Result ChangelogMergeTreeRewriter::RewriteOrProduceChangelog( std::move(merge_function_wrapper), key_comparator, cancellation_checker, drop_delete, produce_data, produce_changelog); auto producer_and_consumer = - std::make_shared>( + std::make_unique>( std::move(producer), create_consumer, /*consumer_thread_num=*/1); - reader_holders.emplace_back(producer_and_consumer); while (true) { if (IsCancelled()) { diff --git a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp index 5ca57f91b..1a818dd01 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp @@ -21,7 +21,6 @@ #include #include "paimon/common/table/special_fields.h" -#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/mergetree/compact/first_row_merge_function_wrapper.h" #include "paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" @@ -85,8 +84,7 @@ LookupMergeTreeCompactRewriter::Create( read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr internal_context, - InternalReadContext::Create(read_context, table_schema, options.ToMap(), - GetSharedArrowPool(read_context->GetMemoryPool()))); + InternalReadContext::Create(read_context, table_schema, options.ToMap())); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr path_factory, path_factory_cache->GetOrCreatePathFactory(options.GetFileFormat()->Identifier())); diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp index e81259cbd..7224f21d4 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp @@ -24,7 +24,6 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "paimon/common/table/special_fields.h" -#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/key_value_data_file_writer_factories.h" #include "paimon/core/io/key_value_meta_projection_consumer.h" @@ -85,8 +84,7 @@ Result> MergeTreeCompactRewriter::Crea read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr internal_context, - InternalReadContext::Create(read_context, table_schema, options.ToMap(), - GetSharedArrowPool(read_context->GetMemoryPool()))); + InternalReadContext::Create(read_context, table_schema, options.ToMap())); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr path_factory, path_factory_cache->GetOrCreatePathFactory(options.GetFileFormat()->Identifier())); @@ -205,14 +203,11 @@ MergeTreeCompactRewriter::CreateRawSortMergeReaderForSection( Status MergeTreeCompactRewriter::MergeReadAndWrite( int32_t output_level, bool drop_delete, const std::vector& section, const MergeTreeCompactRewriter::KeyValueConsumerCreator& create_consumer, - MergeTreeCompactRewriter::KeyValueRollingFileWriter* rolling_writer, - std::vector>* - reader_holders_ptr) { + MergeTreeCompactRewriter::KeyValueRollingFileWriter* rolling_writer) { if (!merge_file_split_read_) { return Status::Invalid( "merge_file_split_read in MergeTreeCompactRewriter cannot be nullptr"); } - auto& reader_holders = *reader_holders_ptr; // prepare sort merge reader PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, CreateDataFilePathFactory(options_.GetFileFormat()->Identifier())); @@ -250,9 +245,8 @@ Status MergeTreeCompactRewriter::MergeReadAndWrite( std::make_unique(std::move(sort_merge_reader), options_.GetWriteBatchSize()); auto async_key_value_producer_consumer = - std::make_shared>( + std::make_unique>( std::move(producer), create_consumer, /*projection_thread_num=*/1); - reader_holders.push_back(async_key_value_producer_consumer); // read KeyValueBatch from SortMergeReader and write to RollingWriter while (true) { if (cancellation_controller_->IsCancelled()) { @@ -273,21 +267,17 @@ Result MergeTreeCompactRewriter::RewriteCompaction( PAIMON_ASSIGN_OR_RAISE(MergeTreeCompactRewriter::KeyValueConsumerCreator create_consumer, GenerateKeyValueConsumer()); auto before = ExtractFilesFromSections(sections); - std::vector> reader_holders; PAIMON_ASSIGN_OR_RAISE(std::unique_ptr rolling_writer, CreateRollingRowWriter(output_level)); ScopeGuard write_guard([&]() -> void { rolling_writer->Abort(); - for (const auto& reader : reader_holders) { - reader->Close(); - } merge_file_split_read_.reset(); }); for (const auto& section : sections) { PAIMON_RETURN_NOT_OK(MergeReadAndWrite(output_level, drop_delete, section, create_consumer, - rolling_writer.get(), &reader_holders)); + rolling_writer.get())); } PAIMON_RETURN_NOT_OK(rolling_writer->Close()); diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.h b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.h index c15e16fc5..513987ffb 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.h +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.h @@ -86,7 +86,6 @@ class MergeTreeCompactRewriter : public CompactRewriter { using KeyValueRollingFileWriter = RollingFileWriter>; - using KeyValueMergeReader = AsyncKeyValueProducerAndConsumer; using KeyValueConsumerCreator = AsyncKeyValueProducerAndConsumer::ConsumerCreator; @@ -99,8 +98,7 @@ class MergeTreeCompactRewriter : public CompactRewriter { Status MergeReadAndWrite(int32_t output_level, bool drop_delete, const std::vector& section, const KeyValueConsumerCreator& create_consumer, - KeyValueRollingFileWriter* rolling_writer, - std::vector>* reader_holders_ptr); + KeyValueRollingFileWriter* rolling_writer); Result> CreateRawSortMergeReaderForSection( const std::vector& section); diff --git a/src/paimon/core/mergetree/lookup_levels.cpp b/src/paimon/core/mergetree/lookup_levels.cpp index 7743d643c..512f28728 100644 --- a/src/paimon/core/mergetree/lookup_levels.cpp +++ b/src/paimon/core/mergetree/lookup_levels.cpp @@ -20,7 +20,6 @@ #include "fmt/format.h" #include "paimon/common/table/special_fields.h" -#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/key_value_data_file_record_reader.h" #include "paimon/core/mergetree/lookup/file_position.h" @@ -69,8 +68,7 @@ Result>> LookupLevels::Create( read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr internal_read_context, - InternalReadContext::Create(read_context, table_schema, options.ToMap(), - GetSharedArrowPool(read_context->GetMemoryPool()))); + InternalReadContext::Create(read_context, table_schema, options.ToMap())); auto split_read = std::make_unique(path_factory, internal_read_context, pool, CreateDefaultExecutor()); diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index 2f46f47ad..3effa5c18 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -161,16 +161,15 @@ Result> AbstractSplitRead::CreateFileBatchReade // TODO(xinyu.lxy): test format table for mosaic format if (context_->EnablePrefetch() && file_format_identifier != "blob" && file_format_identifier != "avro" && file_format_identifier != "mosaic") { - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr prefetch_reader, - PrefetchFileBatchReaderImpl::Create( - data_file_path, data_file_size, reader_builder.get(), options_.GetFileSystem(), - context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), - context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), - executor_, - /*initialize_read_ranges=*/false, context_->ReadAheadCacheEnabled(), - context_->GetCacheConfig(), options_.PrefetchIoMetricsEnabled(), pool_, - arrow_pool_)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr prefetch_reader, + PrefetchFileBatchReaderImpl::Create( + data_file_path, data_file_size, reader_builder.get(), + options_.GetFileSystem(), context_->GetPrefetchMaxParallelNum(), + options_.GetReadBatchSize(), context_->GetPrefetchBatchCount(), + options_.EnableAdaptivePrefetchStrategy(), executor_, + /*initialize_read_ranges=*/false, + context_->ReadAheadCacheEnabled(), context_->GetCacheConfig(), + options_.PrefetchIoMetricsEnabled(), pool_, arrow_pool_)); return std::make_unique(std::move(prefetch_reader)); } else { PAIMON_ASSIGN_OR_RAISE( diff --git a/src/paimon/core/operation/append_only_file_store_write.cpp b/src/paimon/core/operation/append_only_file_store_write.cpp index 9a938f630..f660093a2 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -30,7 +30,6 @@ #include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/arrow_utils.h" -#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/append/append_only_writer.h" #include "paimon/core/append/bucketed_append_compact_manager.h" @@ -291,10 +290,8 @@ Result> AppendOnlyFileStoreWrite::CreateFilesReader .WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, context_builder.Finish()); std::map options = options_.ToMap(); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr internal_read_context, - InternalReadContext::Create(read_context, table_schema_, options, - GetSharedArrowPool(read_context->GetMemoryPool()))); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_read_context, + InternalReadContext::Create(read_context, table_schema_, options)); auto read = std::make_unique(file_store_path_factory_, internal_read_context, pool_, compact_executor_); diff --git a/src/paimon/core/operation/internal_read_context.cpp b/src/paimon/core/operation/internal_read_context.cpp index e686c88ee..3858929e7 100644 --- a/src/paimon/core/operation/internal_read_context.cpp +++ b/src/paimon/core/operation/internal_read_context.cpp @@ -200,8 +200,7 @@ std::optional InternalReadContext::TryResolveSpecialFieldByName( Result> InternalReadContext::Create( const std::shared_ptr& context, const std::shared_ptr& table_schema, - const std::map& options, - const std::shared_ptr& arrow_pool) { + const std::map& options) { PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options, context->GetSpecificFileSystem(), context->GetFileSystemSchemeToIdentifierMap())); @@ -283,9 +282,10 @@ Result> InternalReadContext::Create( PredicateValidator::ValidatePredicateWithLiterals(context->GetPredicate())); } - if (!arrow_pool) { - return Status::Invalid("arrow memory pool is null"); + if (!context->GetMemoryPool()) { + return Status::Invalid("memory pool is null"); } + std::shared_ptr arrow_pool = GetSharedArrowPool(context->GetMemoryPool()); return std::unique_ptr( new InternalReadContext(context, table_schema, read_schema, core_options, arrow_pool)); } diff --git a/src/paimon/core/operation/internal_read_context.h b/src/paimon/core/operation/internal_read_context.h index 69c58236d..cadef31c0 100644 --- a/src/paimon/core/operation/internal_read_context.h +++ b/src/paimon/core/operation/internal_read_context.h @@ -46,8 +46,7 @@ class InternalReadContext { static Result> Create( const std::shared_ptr& read_context, const std::shared_ptr& table_schema, - const std::map& options, - const std::shared_ptr& arrow_pool); + const std::map& options); const CoreOptions& GetCoreOptions() const { return options_; diff --git a/src/paimon/core/operation/internal_read_context_test.cpp b/src/paimon/core/operation/internal_read_context_test.cpp index 7162da6fe..d74a659cc 100644 --- a/src/paimon/core/operation/internal_read_context_test.cpp +++ b/src/paimon/core/operation/internal_read_context_test.cpp @@ -25,7 +25,6 @@ #include "gtest/gtest.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" -#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/data/shredding/map_shared_shredding_schema_utils.h" #include "paimon/defs.h" @@ -40,8 +39,7 @@ Result> CreateInternalReadContextForTest( const std::shared_ptr& read_context, const std::shared_ptr& table_schema, const std::map& options) { - return InternalReadContext::Create(read_context, table_schema, options, - GetSharedArrowPool(read_context->GetMemoryPool())); + return InternalReadContext::Create(read_context, table_schema, options); } } // namespace diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index 4b7a7fab2..02e4387d9 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -94,10 +94,8 @@ class MergeFileSplitReadTest : public ::testing::Test, const std::shared_ptr& read_context, int32_t schema_id = 0) { SchemaManager schema_manager(fs_, read_context->GetPath()); EXPECT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(schema_id)); - EXPECT_OK_AND_ASSIGN( - auto context, - InternalReadContext::Create(read_context, table_schema, read_context->GetOptions(), - GetSharedArrowPool(read_context->GetMemoryPool()))); + EXPECT_OK_AND_ASSIGN(auto context, InternalReadContext::Create(read_context, table_schema, + read_context->GetOptions())); return context; } @@ -365,12 +363,12 @@ class MergeFileSplitReadTest : public ::testing::Test, split_read->CreateReader(split)); batch_readers.emplace_back(std::move(reader)); } - return std::make_unique(std::move(batch_readers), arrow_pool_); + return std::make_unique(std::move(batch_readers), + GetSharedArrowPool(pool_)); } private: std::shared_ptr pool_ = GetDefaultPool(); - std::shared_ptr arrow_pool_ = GetSharedArrowPool(pool_); std::shared_ptr fs_ = std::make_shared(); std::shared_ptr executor_ = CreateDefaultExecutor(); }; diff --git a/src/paimon/core/operation/raw_file_split_read_test.cpp b/src/paimon/core/operation/raw_file_split_read_test.cpp index e6c2184e9..82b132991 100644 --- a/src/paimon/core/operation/raw_file_split_read_test.cpp +++ b/src/paimon/core/operation/raw_file_split_read_test.cpp @@ -144,8 +144,7 @@ class RawFileSplitReadTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN( auto internal_context, - InternalReadContext::Create(read_context, table_schema, table_schema->Options(), - GetSharedArrowPool(read_context->GetMemoryPool()))); + InternalReadContext::Create(read_context, table_schema, table_schema->Options())); auto data_splits = PrepareDataSplits(); const auto& core_options = internal_context->GetCoreOptions(); auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); @@ -174,8 +173,8 @@ class RawFileSplitReadTest : public ::testing::Test { split_read->CreateReader(split)); batch_readers.emplace_back(std::move(reader)); } - auto batch_reader = - std::make_unique(std::move(batch_readers), arrow_pool_); + auto batch_reader = std::make_unique(std::move(batch_readers), + GetSharedArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_TRUE(result_array->Equals(expected_array)); @@ -183,7 +182,6 @@ class RawFileSplitReadTest : public ::testing::Test { private: std::shared_ptr pool_ = GetDefaultPool(); - std::shared_ptr arrow_pool_ = GetSharedArrowPool(pool_); }; // test simple, recall all columns with sequence in data @@ -393,8 +391,7 @@ TEST_F(RawFileSplitReadTest, TestEmptyPlan) { ASSERT_OK_AND_ASSIGN( auto internal_context, - InternalReadContext::Create(read_context, table_schema, table_schema->Options(), - GetSharedArrowPool(read_context->GetMemoryPool()))); + InternalReadContext::Create(read_context, table_schema, table_schema->Options())); const auto& core_options = internal_context->GetCoreOptions(); auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); ASSERT_OK_AND_ASSIGN(std::vector external_paths, @@ -427,7 +424,8 @@ TEST_F(RawFileSplitReadTest, TestEmptyPlan) { std::vector> batch_readers; ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, split_read->CreateReader(data_split)); batch_readers.push_back(std::move(reader)); - auto batch_reader = std::make_unique(std::move(batch_readers), arrow_pool_); + auto batch_reader = + std::make_unique(std::move(batch_readers), GetSharedArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_EQ(result_array, nullptr); @@ -443,8 +441,7 @@ TEST_F(RawFileSplitReadTest, TestMatch) { ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN( std::shared_ptr internal_context, - InternalReadContext::Create(read_context, table_schema, table_schema->Options(), - GetSharedArrowPool(read_context->GetMemoryPool()))); + InternalReadContext::Create(read_context, table_schema, table_schema->Options())); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(/*thread_count=*/2)); auto split_read = std::make_unique( diff --git a/src/paimon/core/table/source/pk_count_reader_test.cpp b/src/paimon/core/table/source/pk_count_reader_test.cpp index 07f677afd..39c0b6b50 100644 --- a/src/paimon/core/table/source/pk_count_reader_test.cpp +++ b/src/paimon/core/table/source/pk_count_reader_test.cpp @@ -26,7 +26,6 @@ #include "gtest/gtest.h" #include "paimon/common/types/data_field.h" -#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/table/source/data_split_impl.h" @@ -79,8 +78,7 @@ class PKCountReaderTest : public testing::Test { PAIMON_ASSIGN_OR_RAISE(auto table_schema, schema_manager.ReadSchema(0)); PAIMON_ASSIGN_OR_RAISE( auto internal_context, - InternalReadContext::Create(read_context, table_schema, table_schema->Options(), - GetSharedArrowPool(read_context->GetMemoryPool()))); + InternalReadContext::Create(read_context, table_schema, table_schema->Options())); return std::shared_ptr(std::move(internal_context)); } diff --git a/src/paimon/core/table/source/table_read.cpp b/src/paimon/core/table/source/table_read.cpp index ef6d80068..be7c6ee05 100644 --- a/src/paimon/core/table/source/table_read.cpp +++ b/src/paimon/core/table/source/table_read.cpp @@ -27,7 +27,6 @@ #include "fmt/format.h" #include "paimon/common/types/data_field.h" -#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/internal_read_context.h" @@ -52,8 +51,7 @@ class MemoryPool; namespace { Result> CreateInternalReadContext( - const std::shared_ptr& context, const std::string& branch, - const std::shared_ptr& arrow_pool) { + const std::shared_ptr& context, const std::string& branch) { std::map tmp_options = context->GetOptions(); std::shared_ptr table_schema; const auto& specific_table_schema = context->GetSpecificTableSchema(); @@ -83,7 +81,7 @@ Result> CreateInternalReadContext( if (branch != BranchManager::DEFAULT_MAIN_BRANCH) { options[Options::BRANCH] = branch; } - return InternalReadContext::Create(context, table_schema, options, arrow_pool); + return InternalReadContext::Create(context, table_schema, options); } Result> CreateTableRead( @@ -117,9 +115,8 @@ Result> CreateTableRead( } Result> NewDataTableRead(const std::shared_ptr& context) { - std::shared_ptr arrow_pool = GetSharedArrowPool(context->GetMemoryPool()); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_context, - CreateInternalReadContext(context, context->GetBranch(), arrow_pool)); + CreateInternalReadContext(context, context->GetBranch())); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr table_read, @@ -134,11 +131,12 @@ Result> NewDataTableRead(const std::shared_ptr fallback_context, - CreateInternalReadContext(context, /*branch=*/scan_fallback_branch.value(), arrow_pool)); + CreateInternalReadContext(context, /*branch=*/scan_fallback_branch.value())); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr fallback_table_read, CreateTableRead(fallback_context, context->GetMemoryPool(), context->GetExecutor())); + const std::shared_ptr& arrow_pool = internal_context->GetArrowMemoryPool(); return std::make_unique(std::move(table_read), std::move(fallback_table_read), arrow_pool); } diff --git a/src/paimon/format/avro/avro_format_writer_test.cpp b/src/paimon/format/avro/avro_format_writer_test.cpp index f89d6b7bb..099d532c8 100644 --- a/src/paimon/format/avro/avro_format_writer_test.cpp +++ b/src/paimon/format/avro/avro_format_writer_test.cpp @@ -52,7 +52,6 @@ class AvroFormatWriterTest : public ::testing::Test { ASSERT_TRUE(dir_); fs_ = std::make_shared(); pool_ = GetDefaultPool(); - arrow_pool_ = GetArrowPool(pool_); } void TearDown() override {} @@ -120,8 +119,9 @@ class AvroFormatWriterTest : public ::testing::Test { void CheckResult(const std::string& file_path, int32_t row_count) const { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs_->Open(file_path)); - ASSERT_OK_AND_ASSIGN(auto file_reader, - AvroFileBatchReader::Create(input_stream, 1024, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN( + auto file_reader, + AvroFileBatchReader::Create(input_stream, 1024, pool_, GetSharedArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(uint64_t num_rows, file_reader->GetNumberOfRows()); ASSERT_EQ(num_rows, row_count); @@ -158,7 +158,6 @@ class AvroFormatWriterTest : public ::testing::Test { std::unique_ptr dir_; std::shared_ptr fs_; std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_F(AvroFormatWriterTest, TestWriteWithVariousBatchSize) { diff --git a/src/paimon/format/blob/blob_file_batch_reader_test.cpp b/src/paimon/format/blob/blob_file_batch_reader_test.cpp index f0a200bda..3b66b6ee8 100644 --- a/src/paimon/format/blob/blob_file_batch_reader_test.cpp +++ b/src/paimon/format/blob/blob_file_batch_reader_test.cpp @@ -36,7 +36,6 @@ class BlobFileBatchReaderTest : public testing::Test, public ::testing::WithPara public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); } void CheckResult(const std::string& table_path, const std::string& paimon_blob_file, @@ -48,10 +47,10 @@ class BlobFileBatchReaderTest : public testing::Test, public ::testing::WithPara std::shared_ptr fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/" + paimon_blob_file)); - ASSERT_OK_AND_ASSIGN( - auto reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( + input_stream, /*batch_size=*/1024, blob_as_descriptor, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, selection_bitmap)); ASSERT_OK_AND_ASSIGN(auto chunked_array, paimon::test::ReadResultCollector::CollectResult(std::move(reader))); @@ -105,7 +104,6 @@ class BlobFileBatchReaderTest : public testing::Test, public ::testing::WithPara private: std::string blob_field_name_; std::shared_ptr pool_; - std::shared_ptr arrow_pool_; }; TEST_P(BlobFileBatchReaderTest, TestSimple) { @@ -168,10 +166,11 @@ TEST_F(BlobFileBatchReaderTest, TestRowNumbers) { ASSERT_OK_AND_ASSIGN( std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); - ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( - input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, std::nullopt)); ASSERT_OK_AND_ASSIGN(auto number_of_rows, reader->GetNumberOfRows()); @@ -210,10 +209,11 @@ TEST_F(BlobFileBatchReaderTest, TestRowNumbersWithSelectionBitmap) { ASSERT_OK_AND_ASSIGN( std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); - ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( - input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); RoaringBitmap32 selection; selection.Add(0); @@ -249,22 +249,24 @@ TEST_F(BlobFileBatchReaderTest, InvalidScenario) { ASSERT_NOK_WITH_MSG( BlobFileBatchReader::Create(input_stream, /*batch_size=*/0, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_), + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_)), "blob file batch reader create failed: read batch size '0' should be larger than zero"); } { ASSERT_NOK_WITH_MSG( BlobFileBatchReader::Create(/*input_stream=*/nullptr, /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_), + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_)), "blob file batch reader create failed: input stream is nullptr"); } { ASSERT_OK_AND_ASSIGN( - auto reader, - BlobFileBatchReader::Create(/*input_stream=*/input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + auto reader, BlobFileBatchReader::Create(/*input_stream=*/input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->GetFileSchema(), "blob file has no self-describing file schema"); ASSERT_TRUE(reader->GetReaderMetrics()); @@ -298,10 +300,11 @@ TEST_P(BlobFileBatchReaderTest, EmptyFile) { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system->Open(dir->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( - input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, std::nullopt)); ASSERT_OK_AND_ASSIGN(auto number_of_rows, reader->GetNumberOfRows()); @@ -324,10 +327,10 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); ASSERT_OK_AND_ASSIGN( - auto reader, - BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + auto reader, BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(/*read_schema=*/nullptr, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "SetReadSchema failed: read schema cannot be nullptr"); @@ -350,10 +353,10 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); ASSERT_OK_AND_ASSIGN( - auto reader, - BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + auto reader, BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "read schema field number 2 is not 1"); @@ -376,10 +379,10 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); ASSERT_OK_AND_ASSIGN( - auto reader, - BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + auto reader, BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "field my_blob_field: large_binary is not BLOB"); @@ -400,10 +403,10 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/data-d7816e8e-6c6d-4e28-9137-837cdf706350-1.blob")); ASSERT_OK_AND_ASSIGN( - auto reader, - BlobFileBatchReader::Create(input_stream, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + auto reader, BlobFileBatchReader::Create(input_stream, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); RoaringBitmap32 roaring; roaring.Add(0); roaring.Add(1); diff --git a/src/paimon/format/blob/blob_format_writer_test.cpp b/src/paimon/format/blob/blob_format_writer_test.cpp index b0afc7ace..5341900db 100644 --- a/src/paimon/format/blob/blob_format_writer_test.cpp +++ b/src/paimon/format/blob/blob_format_writer_test.cpp @@ -116,7 +116,6 @@ class BlobFormatWriterTestBase : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); dir_ = paimon::test::UniqueTestDirectory::Create(); ASSERT_TRUE(dir_); file_system_ = std::make_shared(); @@ -174,11 +173,11 @@ class BlobFormatWriterTestBase : public ::testing::Test { Result> ReadBackAsData() const { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/false, + /*emit_placeholder_sentinel=*/false, + pool_, GetSharedArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); @@ -193,7 +192,6 @@ class BlobFormatWriterTestBase : public ::testing::Test { protected: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; std::unique_ptr dir_; std::shared_ptr output_stream_; std::shared_ptr file_system_; @@ -270,7 +268,8 @@ TEST_P(BlobFormatWriterTest, TestSimple) { ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -467,7 +466,8 @@ TEST_P(BlobFormatWriterTest, TestLargeBlob) { ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -517,7 +517,8 @@ TEST_P(BlobFormatWriterTest, TestAddBatchWithNullValues) { ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -1041,11 +1042,11 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/false, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, @@ -1057,11 +1058,11 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/true, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/false, + /*emit_placeholder_sentinel=*/true, pool_, + GetSharedArrowPool(pool_))); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, @@ -1081,11 +1082,11 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/true, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/true, pool_, + GetSharedArrowPool(pool_))); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, @@ -1112,11 +1113,11 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderWithSelectionBitmap) ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/true, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/false, + /*emit_placeholder_sentinel=*/true, pool_, + GetSharedArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -1145,11 +1146,11 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestSentinelBytesVerbatimWithoutPlacehol ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, - /*blob_as_descriptor=*/false, - /*emit_placeholder_sentinel=*/false, pool_, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, + /*blob_as_descriptor=*/false, + /*emit_placeholder_sentinel=*/false, pool_, + GetSharedArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -1186,7 +1187,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestSentinelPrefixedValueVerbatimInPlace std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, emit_placeholder_sentinel, - pool_, arrow_pool_)); + pool_, GetSharedArrowPool(pool_))); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, diff --git a/src/paimon/format/orc/orc_reader_builder.h b/src/paimon/format/orc/orc_reader_builder.h index f32055109..f1669acbc 100644 --- a/src/paimon/format/orc/orc_reader_builder.h +++ b/src/paimon/format/orc/orc_reader_builder.h @@ -34,12 +34,10 @@ class OrcReaderBuilder : public ReaderBuilder { public: OrcReaderBuilder(const std::map& options, int32_t batch_size) : batch_size_(batch_size), - pool_(GetDefaultPool()), - read_memory_(std::make_shared(pool_)), + read_memory_(std::make_shared(GetDefaultPool())), options_(options) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { - pool_ = pool; read_memory_ = std::make_shared(pool); return this; } @@ -61,7 +59,6 @@ class OrcReaderBuilder : public ReaderBuilder { private: int32_t batch_size_ = -1; - std::shared_ptr pool_; std::shared_ptr read_memory_; std::map options_; }; diff --git a/src/paimon/format/parquet/parquet_vector_io_test.cpp b/src/paimon/format/parquet/parquet_vector_io_test.cpp index db26f1b47..3e6ca2cb4 100644 --- a/src/paimon/format/parquet/parquet_vector_io_test.cpp +++ b/src/paimon/format/parquet/parquet_vector_io_test.cpp @@ -56,7 +56,6 @@ class ParquetVectorIoTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - arrow_pool_ = GetArrowPool(pool_); dir_ = paimon::test::UniqueTestDirectory::Create(); ASSERT_TRUE(dir_); fs_ = dir_->GetFileSystem(); @@ -108,11 +107,12 @@ class ParquetVectorIoTest : public ::testing::Test { fs_->Create(file_path, /*overwrite=*/false)); ::parquet::WriterProperties::Builder properties_builder; properties_builder.max_row_group_length(max_row_group_length); + std::shared_ptr arrow_pool = GetSharedArrowPool(pool_); ASSERT_OK_AND_ASSIGN( std::unique_ptr writer, ParquetFormatWriter::Create(out, arrow::schema(write_type->fields()), properties_builder.build(), - DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_)); + DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool)); ASSERT_OK(writer->AddBatch(c_array.get())); ASSERT_OK(writer->Finish()); ASSERT_OK(out->Close()); @@ -138,10 +138,11 @@ class ParquetVectorIoTest : public ::testing::Test { fs_->Create(file_path, /*overwrite=*/false)); auto arrow_out = std::make_shared(out); ::parquet::WriterProperties::Builder properties_builder; + std::shared_ptr arrow_pool = GetSharedArrowPool(pool_); std::shared_ptr<::parquet::ArrowWriterProperties> arrow_properties = ::parquet::ArrowWriterProperties::Builder().store_schema()->build(); arrow::Status status = ::parquet::arrow::WriteTable( - *std::move(table_result).ValueOrDie(), arrow_pool_.get(), arrow_out, + *std::move(table_result).ValueOrDie(), arrow_pool.get(), arrow_out, /*chunk_size=*/1024, properties_builder.build(), arrow_properties); ASSERT_TRUE(status.ok()) << status.ToString(); ASSERT_OK(out->Close()); @@ -151,12 +152,13 @@ class ParquetVectorIoTest : public ::testing::Test { std::shared_ptr* file_type_out) { ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path)); ASSERT_OK_AND_ASSIGN(int64_t length, in->Length()); - auto in_stream = std::make_shared(in, length, arrow_pool_); + std::shared_ptr arrow_pool = GetSharedArrowPool(pool_); + auto in_stream = std::make_shared(in, length, arrow_pool); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, ParquetFileBatchReader::Create(std::move(in_stream), /*options=*/{}, /*batch_size=*/10, /*file_metadata=*/nullptr, - /*storage_read_bytes=*/nullptr, arrow_pool_, + /*storage_read_bytes=*/nullptr, arrow_pool, /*hints=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(std::unique_ptr c_file_schema, reader->GetFileSchema()); arrow::Result> file_type_result = @@ -173,15 +175,16 @@ class ParquetVectorIoTest : public ::testing::Test { std::unique_ptr* vector_reader_out) { ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path)); ASSERT_OK_AND_ASSIGN(int64_t length, in->Length()); - auto in_stream = std::make_shared(in, length, arrow_pool_); + std::shared_ptr arrow_pool = GetSharedArrowPool(pool_); + auto in_stream = std::make_shared(in, length, arrow_pool); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, ParquetFileBatchReader::Create(std::move(in_stream), options, batch_size, /*file_metadata=*/nullptr, - /*storage_read_bytes=*/nullptr, arrow_pool_, + /*storage_read_bytes=*/nullptr, arrow_pool, /*hints=*/std::nullopt)); std::unique_ptr vector_reader = - std::make_unique(std::move(reader), arrow_pool_); + std::make_unique(std::move(reader), arrow_pool); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); ASSERT_OK(vector_reader->SetReadSchema(c_schema.get(), predicate, @@ -246,7 +249,6 @@ class ParquetVectorIoTest : public ::testing::Test { private: std::shared_ptr pool_; - std::shared_ptr arrow_pool_; std::shared_ptr fs_; std::unique_ptr dir_; }; diff --git a/test/inte/blob_table_inte_test.cpp b/test/inte/blob_table_inte_test.cpp index 30464f0b5..b26c278e4 100644 --- a/test/inte/blob_table_inte_test.cpp +++ b/test/inte/blob_table_inte_test.cpp @@ -92,10 +92,6 @@ class RecordBatch; namespace paimon::test { -struct ReadResult { - std::shared_ptr chunked_array; -}; - class BlobTableInteTest : public testing::Test, public ::testing::WithParamInterface { public: void SetUp() override { @@ -290,11 +286,10 @@ class BlobTableInteTest : public testing::Test, public ::testing::WithParamInter } /// Read from table using a pre-scanned plan and return data after the reader is destroyed. - Result ReadTable(const std::string& table_path, - const std::vector& read_schema, - const std::shared_ptr& plan, - const std::shared_ptr& predicate = nullptr, - const std::map& options = {}) const { + Result> ReadTable( + const std::string& table_path, const std::vector& read_schema, + const std::shared_ptr& plan, const std::shared_ptr& predicate = nullptr, + const std::map& options = {}) const { auto splits = plan->Splits(); ReadContextBuilder read_context_builder(table_path); read_context_builder.SetReadFieldNames(read_schema).SetPredicate(predicate); @@ -308,14 +303,14 @@ class BlobTableInteTest : public testing::Test, public ::testing::WithParamInter PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits)); PAIMON_ASSIGN_OR_RAISE(auto read_result, ReadResultCollector::CollectResult(std::move(batch_reader))); - return ReadResult{std::move(read_result)}; + return read_result; } /// Convenience: scan + read in one call. - Result ScanAndReadResult(const std::string& table_path, - const std::vector& read_schema, - const std::shared_ptr& predicate = nullptr, - const std::vector& row_ranges = {}) const { + Result> ScanAndReadResult( + const std::string& table_path, const std::vector& read_schema, + const std::shared_ptr& predicate = nullptr, + const std::vector& row_ranges = {}) const { PAIMON_ASSIGN_OR_RAISE(auto result_plan, ScanTable(table_path, predicate, row_ranges)); return ReadTable(table_path, read_schema, result_plan, predicate); } @@ -345,13 +340,13 @@ class BlobTableInteTest : public testing::Test, public ::testing::WithParamInter ScanAndReadResult(table_path, read_schema, predicate, row_ranges)); if (!expected_array) { - EXPECT_FALSE(scan_read.chunked_array); + EXPECT_FALSE(scan_read); return Status::OK(); } PAIMON_ASSIGN_OR_RAISE(auto expected_with_row_kind, PrependRowKindColumn(expected_array)); auto expected_chunk_array = std::make_shared(expected_with_row_kind); - EXPECT_TRUE(expected_chunk_array->Equals(scan_read.chunked_array)) - << "result:" << scan_read.chunked_array->ToString() << std::endl + EXPECT_TRUE(expected_chunk_array->Equals(scan_read)) + << "result:" << scan_read->ToString() << std::endl << "expected:" << expected_chunk_array->ToString(); return Status::OK(); } @@ -636,8 +631,8 @@ TEST_P(BlobTableInteTest, TestAppendTableWriteWithBlobAsDescriptorTrue) { // read result contains descriptors pointing to paimon internal blob files // resolve descriptors back to raw bytes, then prepend _VALUE_KIND and compare ASSERT_OK_AND_ASSIGN(auto result, ScanAndReadResult(table_path, schema->field_names())); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_OK_AND_ASSIGN(auto resolved, ConvertDescriptorToRawBlob(read_struct, {"blob"})); ASSERT_OK_AND_ASSIGN(auto expected_with_rk, PrependRowKindColumn(raw_array)); @@ -719,8 +714,8 @@ TEST_P(BlobTableInteTest, TestWriteNullOnMissingFile) { {std::vector{"f0", "f1"}, std::vector{"blob"}}); ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_OK_AND_ASSIGN(auto resolved, ConvertDescriptorToRawBlob(read_struct, {"blob"})); @@ -816,8 +811,8 @@ TEST_P(BlobTableInteTest, TestWriteNullOnFetchFailure) { {std::vector{"f0", "f1"}, std::vector{"blob"}}); ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_OK_AND_ASSIGN(auto resolved, ConvertDescriptorToRawBlob(read_struct, {"blob"})); @@ -878,8 +873,8 @@ TEST_P(BlobTableInteTest, TestWriteNullOnFetchFailureCoversMissingFile) { {std::vector{"f0", "f1"}, std::vector{"blob"}}); ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_OK_AND_ASSIGN(auto resolved, ConvertDescriptorToRawBlob(read_struct, {"blob"})); @@ -1274,8 +1269,8 @@ TEST_P(BlobTableInteTest, TestDataEvolutionBlobPartialUpdateFallback) { std::map read_options = {{Options::BLOB_AS_DESCRIPTOR, "true"}}; ASSERT_OK_AND_ASSIGN(auto desc_result, ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr, read_options)); - ASSERT_TRUE(desc_result.chunked_array); - auto desc_concat = arrow::Concatenate(desc_result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(desc_result); + auto desc_concat = arrow::Concatenate(desc_result->chunks()).ValueOrDie(); auto desc_struct = std::dynamic_pointer_cast(desc_concat); ASSERT_TRUE(desc_struct); ASSERT_OK_AND_ASSIGN(auto resolved, ConvertDescriptorToRawBlob(desc_struct, {"b0"})); @@ -1595,8 +1590,8 @@ TEST_P(BlobTableInteTest, TestDataEvolutionBlobPartialUpdateRowTrackingWithSubra ASSERT_OK_AND_ASSIGN(auto scan_read, ScanAndReadResult(table_path, {"b0", "_ROW_ID", "_SEQUENCE_NUMBER"})); - ASSERT_TRUE(scan_read.chunked_array); - auto concat_array = arrow::Concatenate(scan_read.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(scan_read); + auto concat_array = arrow::Concatenate(scan_read->chunks()).ValueOrDie(); auto struct_array = std::dynamic_pointer_cast(concat_array); ASSERT_TRUE(struct_array); ASSERT_EQ(struct_array->length(), 10); @@ -1627,8 +1622,8 @@ TEST_P(BlobTableInteTest, TestDataEvolutionBlobPartialUpdateRowTrackingWithSubra ScanAndReadResult(table_path, {"b0", "_ROW_ID", "_SEQUENCE_NUMBER"}, /*predicate=*/nullptr, /*row_ranges=*/{Range(1, 2), Range(8, 8)})); - ASSERT_TRUE(range_read.chunked_array); - auto range_concat = arrow::Concatenate(range_read.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(range_read); + auto range_concat = arrow::Concatenate(range_read->chunks()).ValueOrDie(); auto range_struct = std::dynamic_pointer_cast(range_concat); ASSERT_TRUE(range_struct); ASSERT_EQ(range_struct->length(), 3); @@ -1677,8 +1672,8 @@ TEST_P(BlobTableInteTest, TestDataEvolutionBlobPartialUpdateAllPlaceholderRowTra ASSERT_OK_AND_ASSIGN(auto scan_read, ScanAndReadResult(table_path, {"b0", "_ROW_ID", "_SEQUENCE_NUMBER"})); - ASSERT_TRUE(scan_read.chunked_array); - auto concat_array = arrow::Concatenate(scan_read.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(scan_read); + auto concat_array = arrow::Concatenate(scan_read->chunks()).ValueOrDie(); auto struct_array = std::dynamic_pointer_cast(concat_array); ASSERT_TRUE(struct_array); ASSERT_EQ(struct_array->length(), 3); @@ -1757,8 +1752,8 @@ TEST_P(BlobTableInteTest, TestBlobValueEqualToPlaceholderSentinelBytes) { ASSERT_OK(Commit(table_path, commit_msgs1)); ASSERT_OK_AND_ASSIGN(auto scan_read, ScanAndReadResult(table_path, schema->field_names())); - ASSERT_TRUE(scan_read.chunked_array); - auto concat_array = arrow::Concatenate(scan_read.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(scan_read); + auto concat_array = arrow::Concatenate(scan_read->chunks()).ValueOrDie(); auto struct_result = std::dynamic_pointer_cast(concat_array); ASSERT_TRUE(struct_result); ASSERT_EQ(struct_result->length(), 2); @@ -2851,8 +2846,8 @@ TEST_P(BlobTableInteTest, TestBlobDescriptorField) { std::map read_options = {}; ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr, read_options)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_OK_AND_ASSIGN(auto resolved, ConvertDescriptorToRawBlob(read_struct, {"b0", "b1"})); ASSERT_OK_AND_ASSIGN(auto expected_with_rk, PrependRowKindColumn(raw_array)); @@ -2912,8 +2907,8 @@ TEST_P(BlobTableInteTest, TestBlobDescriptorFieldPartialInline) { std::map read_options = {{Options::BLOB_AS_DESCRIPTOR, "true"}}; ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr, read_options)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); // b0,b1 inline descriptor (not repacked), should match input @@ -2991,8 +2986,8 @@ TEST_P(BlobTableInteTest, TestBlobDescriptorMultiCommitAndShuffledReadSchema) { std::map read_options = {{Options::BLOB_AS_DESCRIPTOR, "false"}}; ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, shuffled_read_schema, plan, /*predicate=*/nullptr, read_options)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); // Build expected array in shuffled order from all 3 batches @@ -3026,8 +3021,8 @@ TEST_P(BlobTableInteTest, TestBlobDescriptorMultiCommitAndShuffledReadSchema) { std::map read_options = {{Options::BLOB_AS_DESCRIPTOR, "false"}}; ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, shuffled_read_schema, plan, /*predicate=*/nullptr, read_options)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); // Build expected array in shuffled order from all 3 batches @@ -3399,8 +3394,8 @@ TEST_P(BlobTableInteTest, TestDataEvolutionWithBlobDescriptorField) { std::map read_options = {{Options::BLOB_AS_DESCRIPTOR, "false"}}; ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, read_schema, plan, /*predicate=*/nullptr, read_options)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_EQ(read_struct->length(), 6); @@ -3525,8 +3520,8 @@ TEST_P(BlobTableInteTest, TestBlobViewFieldWithUpstreamTable) { ASSERT_OK_AND_ASSIGN(auto plan, ScanTable(table_path)); ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_EQ(read_struct->length(), 8); ASSERT_OK_AND_ASSIGN(auto result_array, ConvertDescriptorToRawBlob(read_struct, {"view"})); @@ -3558,8 +3553,8 @@ TEST_P(BlobTableInteTest, TestBlobViewFieldWithUpstreamTable) { /*row_ranges=*/{Range(1, 3), Range(5, 5)})); ASSERT_OK_AND_ASSIGN(auto range_result, ReadTable(table_path, schema->field_names(), range_plan, /*predicate=*/nullptr)); - ASSERT_TRUE(range_result.chunked_array); - auto range_concat = arrow::Concatenate(range_result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(range_result); + auto range_concat = arrow::Concatenate(range_result->chunks()).ValueOrDie(); auto range_struct = std::dynamic_pointer_cast(range_concat); ASSERT_EQ(range_struct->length(), 4); ASSERT_OK_AND_ASSIGN(auto range_resolved, @@ -3591,8 +3586,8 @@ TEST_P(BlobTableInteTest, TestBlobViewFieldWithUpstreamTable) { ASSERT_OK_AND_ASSIGN(auto pred_plan, ScanTable(table_path, predicate, /*row_ranges=*/{})); ASSERT_OK_AND_ASSIGN(auto pred_result, ReadTable(table_path, schema->field_names(), pred_plan, predicate)); - ASSERT_TRUE(pred_result.chunked_array); - auto pred_concat = arrow::Concatenate(pred_result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(pred_result); + auto pred_concat = arrow::Concatenate(pred_result->chunks()).ValueOrDie(); auto pred_struct = std::dynamic_pointer_cast(pred_concat); ASSERT_EQ(pred_struct->length(), 8); ASSERT_OK_AND_ASSIGN(auto pred_resolved, ConvertDescriptorToRawBlob(pred_struct, {"view"})); @@ -3706,8 +3701,8 @@ TEST_P(BlobTableInteTest, TestForwardBlobViewReference) { auto source_result, ReadTable(source_table_path, schema->field_names(), source_plan, /*predicate=*/nullptr, {{Options::BLOB_VIEW_RESOLVE_ENABLED, "false"}})); - ASSERT_TRUE(source_result.chunked_array); - auto source_concat = arrow::Concatenate(source_result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(source_result); + auto source_concat = arrow::Concatenate(source_result->chunks()).ValueOrDie(); auto source_struct = std::dynamic_pointer_cast(source_concat); ASSERT_EQ(source_struct->length(), 8); auto forward_f0_array = source_struct->GetFieldByName("f0"); @@ -3734,9 +3729,8 @@ TEST_P(BlobTableInteTest, TestForwardBlobViewReference) { auto raw_target_result, ReadTable(target_table_path, schema->field_names(), target_plan, /*predicate=*/nullptr, {{Options::BLOB_VIEW_RESOLVE_ENABLED, "false"}})); - ASSERT_TRUE(raw_target_result.chunked_array); - auto raw_target_concat = - arrow::Concatenate(raw_target_result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(raw_target_result); + auto raw_target_concat = arrow::Concatenate(raw_target_result->chunks()).ValueOrDie(); auto raw_target_struct = std::dynamic_pointer_cast(raw_target_concat); ASSERT_EQ(raw_target_struct->length(), 8); auto raw_target_view_array = raw_target_struct->GetFieldByName("view"); @@ -3749,8 +3743,8 @@ TEST_P(BlobTableInteTest, TestForwardBlobViewReference) { ASSERT_OK_AND_ASSIGN(auto resolved_result, ReadTable(target_table_path, schema->field_names(), target_plan, /*predicate=*/nullptr)); - ASSERT_TRUE(resolved_result.chunked_array); - auto resolved_concat = arrow::Concatenate(resolved_result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(resolved_result); + auto resolved_concat = arrow::Concatenate(resolved_result->chunks()).ValueOrDie(); auto resolved_struct = std::dynamic_pointer_cast(resolved_concat); ASSERT_EQ(resolved_struct->length(), 8); ASSERT_OK_AND_ASSIGN(auto resolved, ConvertDescriptorToRawBlob(resolved_struct, {"view"})); @@ -3886,8 +3880,8 @@ TEST_P(BlobTableInteTest, TestBlobViewFieldWithUpstreamDescriptorBlob) { ASSERT_OK_AND_ASSIGN(auto plan, ScanTable(table_path)); ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_EQ(read_struct->length(), 4); ASSERT_OK_AND_ASSIGN(auto result_array, ConvertDescriptorToRawBlob(read_struct, {"view"})); @@ -4036,8 +4030,8 @@ TEST_P(BlobTableInteTest, TestBlobViewFieldWithMultipleUpstreamTables) { ASSERT_OK_AND_ASSIGN(auto plan, ScanTable(table_path)); ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, read_fields, plan, /*predicate=*/nullptr)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_EQ(read_struct->length(), 8); ASSERT_OK_AND_ASSIGN(auto result_array, @@ -4116,8 +4110,8 @@ TEST_P(BlobTableInteTest, TestBlobViewSkipsDanglingReferenceOfDeletedRow) { ASSERT_OK_AND_ASSIGN(auto dv_plan, ScanTable(table_path)); ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), dv_plan, /*predicate=*/nullptr)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_EQ(read_struct->length(), 1); @@ -4307,8 +4301,8 @@ TEST_P(BlobTableInteTest, TestBlobViewWithFallbackPath) { ASSERT_OK_AND_ASSIGN(auto plan, ScanTable(table_path)); ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); ASSERT_EQ(read_struct->length(), 2); ASSERT_OK_AND_ASSIGN(auto result_array, ConvertDescriptorToRawBlob(read_struct, {"view"})); @@ -4348,8 +4342,8 @@ TEST_P(BlobTableInteTest, TestReadBlobDescriptorFieldFromJava) { std::map read_options = {{Options::BLOB_AS_DESCRIPTOR, "false"}}; ASSERT_OK_AND_ASSIGN(auto result, ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr, read_options)); - ASSERT_TRUE(result.chunked_array); - auto read_concat = arrow::Concatenate(result.chunked_array->chunks()).ValueOrDie(); + ASSERT_TRUE(result); + auto read_concat = arrow::Concatenate(result->chunks()).ValueOrDie(); auto read_struct = std::dynamic_pointer_cast(read_concat); // After read, b0 and b1 are both descriptor-stored; resolve all back to raw bytes. diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 3b7a19654..a22e3a104 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp @@ -219,13 +219,6 @@ class ReadInteTest : public testing::Test, public ::testing::WithParamInterface< namespace { -struct SystemTableReadResult { - explicit SystemTableReadResult(std::shared_ptr array) - : array(std::move(array)) {} - - std::shared_ptr array; -}; - std::map CollectStringMap( const std::shared_ptr& result) { std::map values; @@ -252,16 +245,17 @@ std::map CollectStringMap( return values; } -std::shared_ptr SingleStructChunk(const SystemTableReadResult& result) { - if (!result.array) { +std::shared_ptr SingleStructChunk( + const std::shared_ptr& result) { + if (!result) { ADD_FAILURE() << "expected non-null system table result"; return nullptr; } - if (result.array->num_chunks() != 1) { - ADD_FAILURE() << "expected one chunk, got " << result.array->num_chunks(); + if (result->num_chunks() != 1) { + ADD_FAILURE() << "expected one chunk, got " << result->num_chunks(); return nullptr; } - auto struct_array = std::dynamic_pointer_cast(result.array->chunk(0)); + auto struct_array = std::dynamic_pointer_cast(result->chunk(0)); if (!struct_array) { ADD_FAILURE() << "expected struct chunk"; } @@ -276,12 +270,10 @@ 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 = {}, - bool read_next_plan = false) { +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) { @@ -312,7 +304,7 @@ Result ReadSystemTable(const std::string& system_table_pa table_read->CreateReader(plan->Splits())); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, ReadResultCollector::CollectResult(std::move(batch_reader))); - return SystemTableReadResult(std::move(result)); + return result; } Status WriteAndFullCompact(std::unique_ptr&& batch, int64_t commit_identifier, @@ -423,10 +415,6 @@ TEST_P(ReadInteTest, TestAppendSimple) { ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); std::shared_ptr read_metrics = batch_reader->GetReaderMetrics(); - std::shared_ptr reader0_metrics = - file_concat_batch_reader->readers_[0]->GetReaderMetrics(); - std::shared_ptr reader1_metrics = - file_concat_batch_reader->readers_[1]->GetReaderMetrics(); auto fields_with_row_kind = read_fields; fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); @@ -443,24 +431,10 @@ TEST_P(ReadInteTest, TestAppendSimple) { ASSERT_TRUE(result_array->Equals(expected_array)); if (param.file_format == "orc") { - ASSERT_OK_AND_ASSIGN(uint64_t reader0_latency, - reader0_metrics->GetCounter("orc.read.inclusive.latency.us")); - ASSERT_OK_AND_ASSIGN(uint64_t reader1_latency, - reader1_metrics->GetCounter("orc.read.inclusive.latency.us")); - uint64_t expected_read_latency = reader0_latency + reader1_latency; - - ASSERT_OK_AND_ASSIGN(uint64_t reader0_io_count, - reader0_metrics->GetCounter("orc.read.io.count")); - ASSERT_OK_AND_ASSIGN(uint64_t reader1_io_count, - reader1_metrics->GetCounter("orc.read.io.count")); - uint64_t expected_read_io_count = reader0_io_count + reader1_io_count; - - ASSERT_OK_AND_ASSIGN(uint64_t result_read_latency, - read_metrics->GetCounter("orc.read.inclusive.latency.us")); - ASSERT_EQ(result_read_latency, expected_read_latency); + ASSERT_OK(read_metrics->GetCounter("orc.read.inclusive.latency.us")); ASSERT_OK_AND_ASSIGN(uint64_t result_read_io_count, read_metrics->GetCounter("orc.read.io.count")); - ASSERT_EQ(result_read_io_count, expected_read_io_count); + ASSERT_GT(result_read_io_count, 0); } }; @@ -923,7 +897,7 @@ TEST(SystemTableReadInteTest, TestReadOptimizedSystemTable) { /*partition_map=*/{}, /*bucket=*/0, {})); ASSERT_OK(WriteAndFullCompact(std::move(batch_1), /*commit_identifier=*/0, helper.get())); - ASSERT_OK_AND_ASSIGN(SystemTableReadResult compacted_result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr compacted_result, ReadSystemTable(table_path + "$ro", options)); std::shared_ptr expected_type = arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), @@ -932,30 +906,28 @@ TEST(SystemTableReadInteTest, TestReadOptimizedSystemTable) { ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( expected_type, {R"([[0, 1, 10], [0, 2, 20]])"}, &expected_compacted) .ok()); - ASSERT_TRUE(compacted_result.array->Equals(expected_compacted)) - << compacted_result.array->ToString(); + ASSERT_TRUE(compacted_result->Equals(expected_compacted)) << compacted_result->ToString(); ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_2, TestHelper::MakeRecordBatch(row_type, R"([[1, 11], [3, 30]])", /*partition_map=*/{}, /*bucket=*/0, {})); ASSERT_OK(helper->WriteAndCommit(std::move(batch_2), /*commit_identifier=*/1, /*expected_commit_messages=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(SystemTableReadResult stale_result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr stale_result, ReadSystemTable(table_path + "$ro", options)); - ASSERT_TRUE(stale_result.array->Equals(expected_compacted)) << stale_result.array->ToString(); + ASSERT_TRUE(stale_result->Equals(expected_compacted)) << stale_result->ToString(); ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_3, TestHelper::MakeRecordBatch(row_type, R"([[2, 21], [3, 31]])", /*partition_map=*/{}, /*bucket=*/0, {})); ASSERT_OK(WriteAndFullCompact(std::move(batch_3), /*commit_identifier=*/2, helper.get())); - ASSERT_OK_AND_ASSIGN(SystemTableReadResult refreshed_result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr refreshed_result, ReadSystemTable(table_path + "$ro", options)); std::shared_ptr expected_refreshed; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( expected_type, {R"([[0, 1, 11], [0, 2, 21], [0, 3, 31]])"}, &expected_refreshed) .ok()); - ASSERT_TRUE(refreshed_result.array->Equals(expected_refreshed)) - << refreshed_result.array->ToString(); + ASSERT_TRUE(refreshed_result->Equals(expected_refreshed)) << refreshed_result->ToString(); } TEST(SystemTableReadInteTest, TestReadOptimizedAppendOnlySystemTableWithStreamingScan) { @@ -985,7 +957,7 @@ TEST(SystemTableReadInteTest, TestReadOptimizedAppendOnlySystemTableWithStreamin ASSERT_OK(helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, /*expected_commit_messages=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(SystemTableReadResult result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadSystemTable(table_path + "$ro", options, /*streaming_mode=*/true)); std::shared_ptr expected_type = arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), @@ -994,7 +966,7 @@ TEST(SystemTableReadInteTest, TestReadOptimizedAppendOnlySystemTableWithStreamin ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( expected_type, {R"([[0, 1, 10], [0, 2, 20]])"}, &expected) .ok()); - ASSERT_TRUE(result.array->Equals(expected)) << result.array->ToString(); + ASSERT_TRUE(result->Equals(expected)) << result->ToString(); } TEST(SystemTableReadInteTest, TestReadOptimizedPrimaryKeyProjectionAndPredicatePushdown) { @@ -1167,7 +1139,7 @@ TEST(SystemTableReadInteTest, TestReadOptimizedSystemTableWithBranch) { /*partition_map=*/{}, /*bucket=*/0, {})); ASSERT_OK(WriteAndFullCompact(std::move(main_batch), /*commit_identifier=*/1, helper.get())); - ASSERT_OK_AND_ASSIGN(SystemTableReadResult result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadSystemTable(table_path + "$branch_rt$ro", options)); std::shared_ptr expected_type = arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), @@ -1176,15 +1148,15 @@ TEST(SystemTableReadInteTest, TestReadOptimizedSystemTableWithBranch) { ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( expected_type, {R"([[0, 1, 10], [0, 2, 20]])"}, &expected) .ok()); - ASSERT_TRUE(result.array->Equals(expected)) << result.array->ToString(); + ASSERT_TRUE(result->Equals(expected)) << result->ToString(); - ASSERT_OK_AND_ASSIGN(SystemTableReadResult main_result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr main_result, ReadSystemTable(table_path + "$ro", options)); std::shared_ptr expected_main; ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( expected_type, {R"([[0, 1, 11], [0, 2, 20], [0, 3, 30]])"}, &expected_main) .ok()); - ASSERT_TRUE(main_result.array->Equals(expected_main)) << main_result.array->ToString(); + ASSERT_TRUE(main_result->Equals(expected_main)) << main_result->ToString(); } TEST(SystemTableReadInteTest, TestReadOptimizedSystemTableWithFirstRowMergeEngine) { @@ -1215,7 +1187,7 @@ TEST(SystemTableReadInteTest, TestReadOptimizedSystemTableWithFirstRowMergeEngin /*partition_map=*/{}, /*bucket=*/0, {})); ASSERT_OK(WriteAndFullCompact(std::move(batch), /*commit_identifier=*/0, helper.get())); - ASSERT_OK_AND_ASSIGN(SystemTableReadResult result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadSystemTable(table_path + "$ro", options)); std::shared_ptr expected_type = arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), @@ -1224,7 +1196,7 @@ TEST(SystemTableReadInteTest, TestReadOptimizedSystemTableWithFirstRowMergeEngin ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( expected_type, {R"([[0, 1, 10], [0, 2, 20]])"}, &expected) .ok()); - ASSERT_TRUE(result.array->Equals(expected)) << result.array->ToString(); + ASSERT_TRUE(result->Equals(expected)) << result->ToString(); } TEST(SystemTableReadInteTest, TestReadFilesSystemTableForPartitionedTable) { @@ -1481,9 +1453,9 @@ TEST(SystemTableReadInteTest, TestReadManifestAndFilesSystemTablesForEmptyTable) catalog->GetTableLocation(Identifier("db1", "tbl1"))); ASSERT_OK_AND_ASSIGN(auto manifests_result, ReadSystemTable(table_path + "$manifests", options)); - ASSERT_EQ(manifests_result.array, nullptr); + ASSERT_EQ(manifests_result, nullptr); ASSERT_OK_AND_ASSIGN(auto files_result, ReadSystemTable(table_path + "$files", options)); - ASSERT_EQ(files_result.array, nullptr); + ASSERT_EQ(files_result, nullptr); } TEST(SystemTableReadInteTest, TestReadTagBranchAndConsumerSystemTables) { @@ -1803,10 +1775,10 @@ TEST(SystemTableReadInteTest, TestStreamingBinlogPacksUpdateBeforeAndAfter) { 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); + ASSERT_TRUE(result); + ASSERT_EQ(result->num_chunks(), 2); auto array = std::dynamic_pointer_cast( - arrow::Concatenate(result.array->chunks()).ValueOrDie()); + arrow::Concatenate(result->chunks()).ValueOrDie()); ASSERT_TRUE(array); ASSERT_EQ(StructFieldNames(array), (std::vector{"rowkind", "pk", "v"})); AssertStructArrayEqualsJson(array, R"([ @@ -4127,7 +4099,7 @@ TEST_P(ReadInteTest, TestSpecificFs) { namespace { -Result ReadGlobalSystemTable( +Result> ReadGlobalSystemTable( const std::string& table_name, Catalog* catalog, const std::shared_ptr& fs, const std::string& warehouse, const std::map& options) { GlobalSystemTableContext ctx; @@ -4161,7 +4133,7 @@ Result ReadGlobalSystemTable( table_read->CreateReader(plan->Splits())); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, ReadResultCollector::CollectResult(std::move(batch_reader))); - return SystemTableReadResult(std::move(result)); + return result; } } // namespace diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 5362f664c..bba491ab4 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -315,10 +315,6 @@ class RealtimeWriteInteTest : public ::testing::Test { protected: using Row = std::tuple; - struct CollectedReadResult { - std::shared_ptr data; - }; - void SetUp() override { pool_ = GetDefaultPool(); dir_ = UniqueTestDirectory::Create("local"); @@ -614,11 +610,10 @@ class RealtimeWriteInteTest : public ::testing::Test { return CreateQueryReader(plan, realtime_context); } - Result ReadPlan(const std::shared_ptr& plan, - const std::shared_ptr& realtime_context, - const std::vector& read_fields, - const std::shared_ptr& predicate, - bool enable_predicate_filter) const { + Result> ReadPlan( + const std::shared_ptr& plan, const std::shared_ptr& realtime_context, + const std::vector& read_fields, const std::shared_ptr& predicate, + bool enable_predicate_filter) const { ReadContextBuilder read_builder(table_path_); read_builder.SetOptions(options_) .SetReadFieldNames(read_fields) @@ -633,7 +628,7 @@ class RealtimeWriteInteTest : public ::testing::Test { table_read->CreateReader(plan->Splits())); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, ReadResultCollector::CollectResult(std::move(reader))); - return CollectedReadResult{std::move(result)}; + return result; } void ReadPlanWithSchemaAndCheck(const std::shared_ptr& plan, @@ -676,11 +671,11 @@ class RealtimeWriteInteTest : public ::testing::Test { Result> ReadRows( const std::shared_ptr& plan, const std::shared_ptr& realtime_context) const { - PAIMON_ASSIGN_OR_RAISE(CollectedReadResult read_result, + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_result, ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, /*predicate=*/nullptr, /*enable_predicate_filter=*/false)); - const std::shared_ptr& result = read_result.data; + const std::shared_ptr& result = read_result; std::vector rows; if (!result) { @@ -2510,7 +2505,7 @@ TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { ASSERT_OK_AND_ASSIGN(std::shared_ptr memory_plan, CreatePlan(realtime_context, scan_predicate)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr memory_result, ReadPlan(memory_plan, realtime_context, read_fields, read_predicate, /*enable_predicate_filter=*/true)); std::shared_ptr expected_memory = @@ -2518,9 +2513,8 @@ TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { [0, "value-2", 2] ])") .ValueOrDie(); - ASSERT_NE(nullptr, memory_result.data); - ASSERT_TRUE( - std::make_shared(expected_memory)->Equals(*memory_result.data)); + ASSERT_NE(nullptr, memory_result); + ASSERT_TRUE(std::make_shared(expected_memory)->Equals(*memory_result)); ASSERT_OK_AND_ASSIGN(std::vector disk_commits, writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); @@ -2532,7 +2526,7 @@ TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { ASSERT_OK_AND_ASSIGN(std::shared_ptr union_plan, CreatePlan(realtime_context, scan_predicate)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult union_result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr union_result, ReadPlan(union_plan, realtime_context, read_fields, read_predicate, /*enable_predicate_filter=*/true)); std::shared_ptr expected_union = @@ -2543,8 +2537,8 @@ TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { [0, "value-5", 5] ])") .ValueOrDie(); - ASSERT_NE(nullptr, union_result.data); - ASSERT_TRUE(std::make_shared(expected_union)->Equals(*union_result.data)); + ASSERT_NE(nullptr, union_result); + ASSERT_TRUE(std::make_shared(expected_union)->Equals(*union_result)); ASSERT_OK(writer->Close()); } @@ -2578,7 +2572,7 @@ TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { ASSERT_OK(writer->Write(std::move(memory_batch))); ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, /*enable_predicate_filter=*/false)); std::shared_ptr result_type = arrow::struct_( @@ -2592,9 +2586,9 @@ TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { [0, 5, "value-5", "p0"] ])") .ValueOrDie(); - ASSERT_NE(nullptr, result.data); - ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) - << result.data->ToString(); + ASSERT_NE(nullptr, result); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) + << result->ToString(); ASSERT_OK(writer->Close()); } @@ -2640,11 +2634,10 @@ TEST_F(RealtimeWriteInteTest, TestMemoryBatchStatisticsPredicatePushdown) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); PAIMON_ASSIGN_OR_RAISE( - CollectedReadResult result, + std::shared_ptr result, ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, /*enable_predicate_filter=*/false)); - std::shared_ptr actual = - result.data ? result.data : make_expected("[]"); + std::shared_ptr actual = result ? result : make_expected("[]"); if (!expected[i]->Equals(*actual)) { return Status::Invalid("unexpected real-time candidate rows: " + actual->ToString()); @@ -2715,7 +2708,7 @@ TEST_F(RealtimeWriteInteTest, TestMemoryBatchStatisticsPredicatePushdownWithDisk PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, Literal(static_cast(3))); ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, /*enable_predicate_filter=*/false)); @@ -2731,9 +2724,9 @@ TEST_F(RealtimeWriteInteTest, TestMemoryBatchStatisticsPredicatePushdownWithDisk [0, 12, "value-12", "p0"] ])") .ValueOrDie(); - ASSERT_NE(nullptr, result.data); - ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) - << result.data->ToString(); + ASSERT_NE(nullptr, result); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) + << result->ToString(); ASSERT_OK(writer->Close()); } @@ -2778,7 +2771,7 @@ TEST_F(RealtimeWriteInteTest, TestNullPredicateForMemoryAndDisk) { std::shared_ptr predicate = PredicateBuilder::IsNull( /*field_index=*/1, /*field_name=*/"payload", FieldType::STRING); ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, /*enable_predicate_filter=*/false)); @@ -2792,9 +2785,9 @@ TEST_F(RealtimeWriteInteTest, TestNullPredicateForMemoryAndDisk) { [0, 5, "memory-value-5", "p0"] ])") .ValueOrDie(); - ASSERT_NE(nullptr, result.data); - ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) - << result.data->ToString(); + ASSERT_NE(nullptr, result); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) + << result->ToString(); ASSERT_OK(writer->Close()); } @@ -2831,7 +2824,7 @@ TEST_F(RealtimeWriteInteTest, TestUnionReadAfterColumnRename) { ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(second_context, /*predicate=*/nullptr)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadPlan(plan, second_context, {"id", "renamed_payload", "pt"}, /*predicate=*/nullptr, /*enable_predicate_filter=*/false)); std::shared_ptr result_type = arrow::struct_( @@ -2846,9 +2839,9 @@ TEST_F(RealtimeWriteInteTest, TestUnionReadAfterColumnRename) { [0, 4, "value-4", "p0"] ])") .ValueOrDie(); - ASSERT_NE(nullptr, result.data); - ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) - << result.data->ToString(); + ASSERT_NE(nullptr, result); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) + << result->ToString(); ASSERT_OK(second_writer->Close()); } From 2975f60dd6850d5698e9305333ca388d20905a6a Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Sat, 29 Aug 2026 23:02:53 +0800 Subject: [PATCH 3/8] fix little --- .../reader/concat_batch_reader_test.cpp | 58 +++++++++++++++++++ .../common/utils/arrow/mem_utils_test.cpp | 2 +- ...pply_deletion_vector_batch_reader_test.cpp | 10 ++-- .../core/io/concat_key_value_record_reader.h | 22 +++++-- .../concat_key_value_record_reader_test.cpp | 34 ++++++++++- .../core/mergetree/compact/loser_tree.cpp | 19 +++++- .../core/mergetree/compact/loser_tree.h | 23 ++++---- .../compact/sort_merge_reader_test.cpp | 39 ++++++++++++- .../sort_merge_reader_with_min_heap.cpp | 25 ++++---- .../compact/sort_merge_reader_with_min_heap.h | 28 +++++---- .../operation/internal_read_context_test.cpp | 52 +++++++---------- .../format/orc/complex_predicate_test.cpp | 8 +-- .../format/orc/orc_file_batch_reader.cpp | 17 ++---- src/paimon/format/orc/orc_file_batch_reader.h | 18 +----- .../format/orc/orc_file_batch_reader_test.cpp | 27 ++++++--- src/paimon/format/orc/orc_format_defs.h | 16 +++++ src/paimon/format/orc/orc_reader_wrapper.cpp | 4 +- src/paimon/format/orc/orc_reader_wrapper.h | 22 ++++--- .../format/orc/orc_reader_wrapper_test.cpp | 19 +++++- .../format/orc/predicate_pushdown_test.cpp | 8 +-- 20 files changed, 309 insertions(+), 142 deletions(-) diff --git a/src/paimon/common/reader/concat_batch_reader_test.cpp b/src/paimon/common/reader/concat_batch_reader_test.cpp index a697e9ad6..654066743 100644 --- a/src/paimon/common/reader/concat_batch_reader_test.cpp +++ b/src/paimon/common/reader/concat_batch_reader_test.cpp @@ -29,6 +29,7 @@ #include "arrow/array/array_nested.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/memory/memory_pool.h" @@ -67,6 +68,36 @@ class LifetimeTrackingBatchReader : public BatchReader { std::shared_ptr lifetime_; }; +class FixedMetricsBatchReader : public BatchReader { + public: + FixedMetricsBatchReader(std::unique_ptr reader, uint64_t latency, + uint64_t io_count) + : reader_(std::move(reader)), metrics_(std::make_shared()) { + metrics_->SetCounter("orc.read.inclusive.latency.us", latency); + metrics_->SetCounter("orc.read.io.count", io_count); + } + + Result NextBatch() override { + return reader_->NextBatch(); + } + + Result NextBatchWithBitmap() override { + return reader_->NextBatchWithBitmap(); + } + + void Close() override { + reader_->Close(); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + + private: + std::unique_ptr reader_; + std::shared_ptr metrics_; +}; + class ConcatBatchReaderTest : public ::testing::Test { void SetUp() override { pool_ = GetDefaultPool(); @@ -225,4 +256,31 @@ TEST_F(ConcatBatchReaderTest, TestReleaseReaderAtEof) { ReaderUtils::ReleaseReadBatch(std::move(batch.first)); } +TEST_F(ConcatBatchReaderTest, TestCollectMetricsAfterReleasingReaders) { + std::shared_ptr type = arrow::struct_({arrow::field("f1", arrow::int32())}); + std::shared_ptr data1 = + arrow::ipc::internal::json::ArrayFromJSON(type, "[[1]]").ValueOrDie(); + std::shared_ptr data2 = + arrow::ipc::internal::json::ArrayFromJSON(type, "[[2]]").ValueOrDie(); + + std::vector> readers; + readers.push_back(std::make_unique( + std::make_unique(data1, type, /*read_batch_size=*/1), + /*latency=*/11, /*io_count=*/2)); + readers.push_back(std::make_unique( + std::make_unique(data2, type, /*read_batch_size=*/1), + /*latency=*/17, /*io_count=*/3)); + auto concat_reader = + std::make_unique(std::move(readers), GetSharedArrowPool(pool_)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadResultCollector::CollectResult(concat_reader.get())); + ASSERT_EQ(result->length(), 2); + std::shared_ptr metrics = concat_reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t latency, metrics->GetCounter("orc.read.inclusive.latency.us")); + ASSERT_EQ(latency, 28); + ASSERT_OK_AND_ASSIGN(uint64_t io_count, metrics->GetCounter("orc.read.io.count")); + ASSERT_EQ(io_count, 5); +} + } // namespace paimon::test diff --git a/src/paimon/common/utils/arrow/mem_utils_test.cpp b/src/paimon/common/utils/arrow/mem_utils_test.cpp index 67878b716..820f3b95a 100644 --- a/src/paimon/common/utils/arrow/mem_utils_test.cpp +++ b/src/paimon/common/utils/arrow/mem_utils_test.cpp @@ -271,7 +271,7 @@ TEST(MemUtilsTest, TestAddArrowArrayLifetimeDeduplicatesSameOwner) { array.private_data = new TrackingArrowArrayPrivateData{nullptr, &release_order}; ASSERT_OK(AddArrowArrayLifetime(&array, lifetime)); - const long use_count = lifetime.use_count(); + const int64_t use_count = lifetime.use_count(); ASSERT_OK(AddArrowArrayLifetime(&array, lifetime)); ASSERT_EQ(use_count, lifetime.use_count()); diff --git a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp index 8cc1d800c..a1baa2868 100644 --- a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp +++ b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp @@ -50,8 +50,6 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, int_type_ = arrow::int32(); target_type_ = arrow::struct_({arrow::field("f1", int_type_)}); - pool_ = GetDefaultPool(); - arrow_pool_ = GetSharedArrowPool(pool_); fs_ = std::make_shared(); ASSERT_OK_AND_ASSIGN(executor_, CreateDefaultExecutor(/*thread_count=*/2)); } @@ -71,13 +69,15 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, void CheckResult(const std::string& data_str, const std::vector& dv_data, const std::string& expected_str) { + std::shared_ptr pool = GetDefaultPool(); + std::shared_ptr arrow_pool = GetSharedArrowPool(pool); auto f1 = arrow::ipc::internal::json::ArrayFromJSON(int_type_, data_str).ValueOrDie(); std::shared_ptr data = arrow::StructArray::Make({f1}, target_type_->fields()).ValueOrDie(); int32_t prefetch_batch_count = 3; for (int32_t batch_size : {1, 2, 4, 10}) { - auto dv = DeletionVector::FromPrimitiveArray(dv_data, pool_.get()); + auto dv = DeletionVector::FromPrimitiveArray(dv_data, pool.get()); std::unique_ptr file_batch_reader; bool enable_prefetch = GetParam(); if (enable_prefetch) { @@ -90,7 +90,7 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, arrow_pool_)); + /*enable_io_metrics=*/false, pool, arrow_pool)); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); @@ -116,8 +116,6 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, private: std::shared_ptr int_type_; std::shared_ptr target_type_; - std::shared_ptr pool_; - std::shared_ptr arrow_pool_; std::shared_ptr fs_; std::shared_ptr executor_; }; diff --git a/src/paimon/core/io/concat_key_value_record_reader.h b/src/paimon/core/io/concat_key_value_record_reader.h index 69a1da6d2..7a3f545ed 100644 --- a/src/paimon/core/io/concat_key_value_record_reader.h +++ b/src/paimon/core/io/concat_key_value_record_reader.h @@ -37,7 +37,7 @@ class ConcatKeyValueRecordReader : public KeyValueRecordReader { public: explicit ConcatKeyValueRecordReader( std::vector>&& readers) - : readers_(std::move(readers)) {} + : finished_reader_metrics_(std::make_shared()), readers_(std::move(readers)) {} Result> NextBatch() override { while (current_ < readers_.size()) { @@ -48,7 +48,7 @@ class ConcatKeyValueRecordReader : public KeyValueRecordReader { return iterator; } // current meets eof, move to next reader - current_reader->Close(); + CloseAndReleaseReader(current_); current_++; } // read finish @@ -57,15 +57,29 @@ class ConcatKeyValueRecordReader : public KeyValueRecordReader { void Close() override { for (; current_ < readers_.size(); current_++) { - readers_[current_]->Close(); + CloseAndReleaseReader(current_); } } std::shared_ptr GetReaderMetrics() const override { - return MetricsImpl::CollectReadMetrics(readers_); + auto metrics = std::make_shared(); + metrics->Merge(finished_reader_metrics_); + metrics->Merge(MetricsImpl::CollectReadMetrics(readers_)); + return metrics; } private: + void CloseAndReleaseReader(size_t reader_index) { + std::unique_ptr& reader = readers_[reader_index]; + if (!reader) { + return; + } + reader->Close(); + finished_reader_metrics_->Merge(reader->GetReaderMetrics()); + reader.reset(); + } + + std::shared_ptr finished_reader_metrics_; std::vector> readers_; size_t current_{0}; }; diff --git a/src/paimon/core/io/concat_key_value_record_reader_test.cpp b/src/paimon/core/io/concat_key_value_record_reader_test.cpp index b9da3cfd3..b9a2ac234 100644 --- a/src/paimon/core/io/concat_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/concat_key_value_record_reader_test.cpp @@ -37,6 +37,30 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { + +class ConcatLifetimeTrackingKeyValueRecordReader : public KeyValueRecordReader { + public: + ConcatLifetimeTrackingKeyValueRecordReader(std::unique_ptr reader, + const std::shared_ptr& lifetime) + : reader_(std::move(reader)), lifetime_(lifetime) {} + + Result> NextBatch() override { + return reader_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + reader_->Close(); + } + + private: + std::unique_ptr reader_; + std::shared_ptr lifetime_; +}; + class ConcatKeyValueRecordReaderTest : public testing::Test { public: void SetUp() override { @@ -49,13 +73,18 @@ class ConcatKeyValueRecordReaderTest : public testing::Test { const std::shared_ptr& value_schema) const { for (auto batch_size : {1, 2, 3, 4, 100}) { std::vector> record_readers; + std::vector> reader_lifetimes; for (const auto& src_array : src_array_vec) { auto src_type = src_array->type(); auto file_batch_reader = std::make_unique( src_array, src_type, /*batch_size=*/batch_size); auto record_reader = std::make_unique( std::move(file_batch_reader), key_schema, value_schema, /*level=*/0, pool_); - record_readers.push_back(std::move(record_reader)); + std::shared_ptr lifetime = std::make_shared(0); + reader_lifetimes.push_back(lifetime); + record_readers.push_back( + std::make_unique( + std::move(record_reader), lifetime)); } auto concat_record_reader = @@ -67,6 +96,9 @@ class ConcatKeyValueRecordReaderTest : public testing::Test { concat_record_reader.get()))); KeyValueChecker::CheckResult(expected, results, key_schema->num_fields(), value_schema->num_fields()); + for (const auto& reader_lifetime : reader_lifetimes) { + ASSERT_TRUE(reader_lifetime.expired()); + } } } diff --git a/src/paimon/core/mergetree/compact/loser_tree.cpp b/src/paimon/core/mergetree/compact/loser_tree.cpp index e93211903..596cbc5b8 100644 --- a/src/paimon/core/mergetree/compact/loser_tree.cpp +++ b/src/paimon/core/mergetree/compact/loser_tree.cpp @@ -21,19 +21,32 @@ #include #include +#include "paimon/common/metrics/metrics_impl.h" + namespace paimon { LoserTree::LoserTree(std::vector>&& readers, const CompareFunc& first_comparator, const CompareFunc& second_comparator) : size_(readers.size()), initialized_(false), - readers_holder_(std::move(readers)), tree_(size_), first_comparator_(first_comparator), second_comparator_(second_comparator) { leaves_.reserve(size_); - for (const auto& reader : readers_holder_) { - leaves_.emplace_back(reader.get()); + for (auto& reader : readers) { + leaves_.emplace_back(std::move(reader)); + } +} + +std::shared_ptr LoserTree::GetReaderMetrics() const { + auto metrics = std::make_shared(); + for (const auto& leaf : leaves_) { + if (leaf.reader) { + metrics->Merge(leaf.reader->GetReaderMetrics()); + } else { + metrics->Merge(leaf.finished_reader_metrics); + } } + return metrics; } Status LoserTree::InitializeIfNeeded() { diff --git a/src/paimon/core/mergetree/compact/loser_tree.h b/src/paimon/core/mergetree/compact/loser_tree.h index 1417d8de8..1bde92131 100644 --- a/src/paimon/core/mergetree/compact/loser_tree.h +++ b/src/paimon/core/mergetree/compact/loser_tree.h @@ -24,7 +24,6 @@ #include #include -#include "paimon/common/metrics/metrics_impl.h" #include "paimon/core/io/key_value_record_reader.h" #include "paimon/core/key_value.h" #include "paimon/result.h" @@ -52,13 +51,13 @@ class LoserTree { /// Peek the current winner, mainly for key comparisons. const std::optional& PeekWinner() const; - std::shared_ptr GetReaderMetrics() const { - return MetricsImpl::CollectReadMetrics(readers_holder_); - } + std::shared_ptr GetReaderMetrics() const; void Close() { - for (const auto& reader : readers_holder_) { - reader->Close(); + for (const auto& leaf : leaves_) { + if (leaf.reader) { + leaf.reader->Close(); + } } } @@ -97,7 +96,8 @@ class LoserTree { } struct LeafIterator { - explicit LeafIterator(KeyValueRecordReader* reader) : reader(reader) {} + explicit LeafIterator(std::unique_ptr&& reader) + : reader(std::move(reader)) {} const std::optional& Peek() const { return kv; @@ -127,7 +127,10 @@ class LoserTree { PAIMON_ASSIGN_OR_RAISE(iterator, reader->NextBatch()); if (!iterator) { // read eof + iterator.reset(); reader->Close(); + finished_reader_metrics = reader->GetReaderMetrics(); + reader.reset(); end_of_input = true; kv = std::nullopt; } else { @@ -148,7 +151,8 @@ class LoserTree { bool end_of_input = false; int32_t first_same_key_index = -1; State state = State::WINNER_WITH_NEW_KEY; - KeyValueRecordReader* reader; + std::unique_ptr reader; + std::shared_ptr finished_reader_metrics; std::unique_ptr iterator; std::optional kv; }; @@ -156,9 +160,6 @@ class LoserTree { private: int32_t size_; bool initialized_; - // must hold all readers, as data array is allocated by the pool of data file - // reader - std::vector> readers_holder_; std::vector tree_; std::vector leaves_; diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp b/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp index dabee55c4..cc7bf5e15 100644 --- a/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp +++ b/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp @@ -52,6 +52,30 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { + +class SortMergeLifetimeTrackingKeyValueRecordReader : public KeyValueRecordReader { + public: + SortMergeLifetimeTrackingKeyValueRecordReader(std::unique_ptr reader, + const std::shared_ptr& lifetime) + : reader_(std::move(reader)), lifetime_(lifetime) {} + + Result> NextBatch() override { + return reader_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + reader_->Close(); + } + + private: + std::unique_ptr reader_; + std::shared_ptr lifetime_; +}; + class SortMergeReaderTest : public testing::Test { public: void SetUp() override { @@ -92,7 +116,7 @@ class SortMergeReaderTest : public testing::Test { const std::shared_ptr& user_defined_seq_comparator, const std::shared_ptr& key_schema, const std::shared_ptr& value_schema, int32_t batch_size, bool need_merge, - bool ignore_delete) const { + bool ignore_delete, std::vector>* reader_lifetimes) const { auto mfunc = std::make_unique(ignore_delete); auto merge_function_wrapper = std::make_shared(std::move(mfunc)); @@ -111,8 +135,13 @@ class SortMergeReaderTest : public testing::Test { std::move(file_batch_reader), key_schema, value_schema, /*level=*/0, pool_); std::vector> readers; readers.push_back(std::move(record_reader)); + std::unique_ptr concat_reader = + std::make_unique(std::move(readers)); + std::shared_ptr lifetime = std::make_shared(0); + reader_lifetimes->push_back(lifetime); concat_readers.push_back( - std::make_unique(std::move(readers))); + std::make_unique( + std::move(concat_reader), lifetime)); } return std::make_unique(std::move(concat_readers), user_key_comparator, @@ -129,15 +158,19 @@ class SortMergeReaderTest : public testing::Test { const std::vector& expected, bool need_merge, bool ignore_delete = false) const { for (auto batch_size : {1, 2, 3, 4, 100}) { + std::vector> reader_lifetimes; auto sort_merge_reader = CreateSortMergeReader( src_array_vec, user_key_comparator, user_defined_seq_comparator, key_schema, - value_schema, batch_size, need_merge, ignore_delete); + value_schema, batch_size, need_merge, ignore_delete, &reader_lifetimes); ASSERT_OK_AND_ASSIGN( std::vector results, (ReadResultCollector::CollectKeyValueResult< SortMergeReader, SortMergeReader::Iterator>(sort_merge_reader.get()))); KeyValueChecker::CheckResult(expected, results, key_schema->num_fields(), value_schema->num_fields()); + for (const auto& reader_lifetime : reader_lifetimes) { + ASSERT_TRUE(reader_lifetime.expired()); + } } } diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp index 08629e638..fc592accf 100644 --- a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp +++ b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp @@ -30,35 +30,40 @@ SortMergeReaderWithMinHeap::SortMergeReaderWithMinHeap( const std::shared_ptr& user_defined_seq_comparator, const std::shared_ptr>& merge_function_wrapper) : need_merge_(merge_function_wrapper != nullptr), - readers_holder_(std::move(readers)), + readers_(std::move(readers)), + finished_reader_metrics_(std::make_shared()), user_key_comparator_(user_key_comparator), merge_function_wrapper_(merge_function_wrapper), min_heap_(HeapSorter(user_key_comparator, user_defined_seq_comparator)) { - next_batch_readers_.reserve(readers_holder_.size()); - for (auto& reader : readers_holder_) { - next_batch_readers_.push_back(reader.get()); + next_batch_reader_indices_.reserve(readers_.size()); + for (size_t reader_index = 0; reader_index < readers_.size(); reader_index++) { + next_batch_reader_indices_.push_back(reader_index); } } Result> SortMergeReaderWithMinHeap::NextBatch() { - for (auto* reader : next_batch_readers_) { + for (size_t reader_index : next_batch_reader_indices_) { + std::unique_ptr& reader = readers_[reader_index]; + assert(reader); while (true) { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, reader->NextBatch()); if (!iterator) { // no more batches, permanently remove this reader reader->Close(); + finished_reader_metrics_->Merge(reader->GetReaderMetrics()); + reader.reset(); break; } PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext()); if (has_next) { PAIMON_ASSIGN_OR_RAISE(KeyValue kv, iterator->Next()); - min_heap_.emplace(std::move(kv), std::move(iterator), reader); + min_heap_.emplace(std::move(kv), std::move(iterator), reader_index); break; } } } - next_batch_readers_.clear(); + next_batch_reader_indices_.clear(); if (min_heap_.empty()) { return std::unique_ptr(); } @@ -85,7 +90,7 @@ Result SortMergeReaderWithMinHeap::Iterator::HasNext() { } Result SortMergeReaderWithMinHeap::Iterator::NextImpl() { - assert(reader_->next_batch_readers_.empty()); + assert(reader_->next_batch_reader_indices_.empty()); // add previously polled elements back to priority queue for (auto& element : reader_->polled_) { PAIMON_ASSIGN_OR_RAISE(bool updated, element.Update()); @@ -94,11 +99,11 @@ Result SortMergeReaderWithMinHeap::Iterator::NextImpl() { reader_->min_heap_.push(std::move(element)); } else { // reach end of batch, clean up - reader_->next_batch_readers_.push_back(std::move(element.reader)); + reader_->next_batch_reader_indices_.push_back(element.reader_index); } } reader_->polled_.clear(); - if (!reader_->next_batch_readers_.empty()) { + if (!reader_->next_batch_reader_indices_.empty()) { return false; } diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h index 3d991004d..667af1242 100644 --- a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h +++ b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h @@ -69,8 +69,10 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { Result> NextBatch() override; void Close() override { - for (const auto& reader : readers_holder_) { - reader->Close(); + for (const auto& reader : readers_) { + if (reader) { + reader->Close(); + } } if (merge_function_wrapper_) { merge_function_wrapper_->Reset(); @@ -78,21 +80,23 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { } std::shared_ptr GetReaderMetrics() const override { - return MetricsImpl::CollectReadMetrics(readers_holder_); + auto metrics = std::make_shared(); + metrics->Merge(finished_reader_metrics_); + metrics->Merge(MetricsImpl::CollectReadMetrics(readers_)); + return metrics; } private: struct Element { Element(KeyValue&& _kv, std::unique_ptr&& _iterator, - KeyValueRecordReader* _reader) - : kv(std::move(_kv)), reader(_reader), iterator(std::move(_iterator)) { + size_t _reader_index) + : kv(std::move(_kv)), reader_index(_reader_index), iterator(std::move(_iterator)) { assert(iterator); - assert(reader); } Element(Element&& other) noexcept : kv(std::move(other.kv)) { iterator = std::move(other.iterator); - reader = other.reader; + reader_index = other.reader_index; } Element& operator=(Element&& other) noexcept { @@ -101,7 +105,7 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { } kv = std::move(other.kv); iterator = std::move(other.iterator); - reader = other.reader; + reader_index = other.reader_index; return *this; } @@ -117,7 +121,7 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { public: KeyValue kv; - KeyValueRecordReader* reader; + size_t reader_index; std::unique_ptr iterator; }; @@ -149,9 +153,9 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { private: const bool need_merge_; - // must hold all readers, as data array is allocated by the pool of data file reader - std::vector> readers_holder_; - std::vector next_batch_readers_; + std::vector> readers_; + std::vector next_batch_reader_indices_; + std::shared_ptr finished_reader_metrics_; std::shared_ptr user_key_comparator_; std::shared_ptr> merge_function_wrapper_; std::priority_queue, HeapSorter> min_heap_; diff --git a/src/paimon/core/operation/internal_read_context_test.cpp b/src/paimon/core/operation/internal_read_context_test.cpp index d74a659cc..371d2d762 100644 --- a/src/paimon/core/operation/internal_read_context_test.cpp +++ b/src/paimon/core/operation/internal_read_context_test.cpp @@ -33,16 +33,6 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { -namespace { - -Result> CreateInternalReadContextForTest( - const std::shared_ptr& read_context, - const std::shared_ptr& table_schema, - const std::map& options) { - return InternalReadContext::Create(read_context, table_schema, options); -} - -} // namespace TEST(InternalReadContext, TestReadWithUnspecifiedSchema) { // no read schema is specified, read all fields @@ -52,8 +42,8 @@ TEST(InternalReadContext, TestReadWithUnspecifiedSchema) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - CreateInternalReadContextForTest(std::move(read_context), table_schema, - table_schema->Options())); + InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(0, arrow::field("f0", arrow::utf8())), DataField(1, arrow::field("f1", arrow::int32())), DataField(2, arrow::field("f2", arrow::int32())), @@ -70,8 +60,8 @@ TEST(InternalReadContext, TestReadWithSpecifiedSchema) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - CreateInternalReadContextForTest(std::move(read_context), table_schema, - table_schema->Options())); + InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8()))}; auto expected_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); @@ -86,8 +76,8 @@ TEST(InternalReadContext, TestReadWithSpecifiedFieldId) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - CreateInternalReadContextForTest(std::move(read_context), table_schema, - table_schema->Options())); + InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8()))}; auto expected_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); @@ -105,8 +95,8 @@ TEST(InternalReadContext, TestReadWithSpecifiedFieldIdAndSchema) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - CreateInternalReadContextForTest(std::move(read_context), table_schema, - table_schema->Options())); + InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8()))}; auto expected_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); @@ -128,7 +118,7 @@ TEST(InternalReadContext, TestReadWithRowTrackingAndScoreFields) { new_options[Options::DATA_EVOLUTION_ENABLED] = "true"; ASSERT_OK_AND_ASSIGN( auto internal_context, - CreateInternalReadContextForTest(std::move(read_context), table_schema, new_options)); + InternalReadContext::Create(std::move(read_context), table_schema, new_options)); std::vector read_fields = { DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8())), SpecialFields::RowId(), @@ -144,8 +134,8 @@ TEST(InternalReadContext, TestReadWithRowTrackingAndScoreFields) { ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); - ASSERT_NOK_WITH_MSG(CreateInternalReadContextForTest(std::move(read_context), table_schema, - table_schema->Options()), + ASSERT_NOK_WITH_MSG(InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options()), "Get field _ROW_ID failed: not exist in table schema"); } { @@ -156,8 +146,8 @@ TEST(InternalReadContext, TestReadWithRowTrackingAndScoreFields) { ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); - ASSERT_NOK_WITH_MSG(CreateInternalReadContextForTest(std::move(read_context), table_schema, - table_schema->Options()), + ASSERT_NOK_WITH_MSG(InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options()), "Get field _INDEX_SCORE failed: not exist in table schema"); } } @@ -170,8 +160,8 @@ TEST(InternalReadContext, TestReadWithValueKindField) { SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); ASSERT_OK_AND_ASSIGN(auto internal_context, - CreateInternalReadContextForTest(std::move(read_context), table_schema, - table_schema->Options())); + InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options())); std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), SpecialFields::ValueKind(), DataField(0, arrow::field("f0", arrow::utf8()))}; @@ -197,7 +187,7 @@ TEST(InternalReadContext, TestReadWithFieldIdsAndSpecialFields) { new_options[Options::DATA_EVOLUTION_ENABLED] = "true"; ASSERT_OK_AND_ASSIGN( auto internal_context, - CreateInternalReadContextForTest(std::move(read_context), table_schema, new_options)); + InternalReadContext::Create(std::move(read_context), table_schema, new_options)); std::vector read_fields = { DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8())), SpecialFields::RowId(), @@ -226,7 +216,7 @@ TEST(InternalReadContext, TestReadWithProjectedSchemaAndSpecialFields) { ASSERT_OK_AND_ASSIGN(auto unique_read_context, context_builder.Finish()); std::shared_ptr read_context = std::move(unique_read_context); ASSERT_NOK_WITH_MSG( - CreateInternalReadContextForTest(read_context, table_schema, table_schema->Options()), + InternalReadContext::Create(read_context, table_schema, table_schema->Options()), "not exist in table schema"); } @@ -245,7 +235,7 @@ TEST(InternalReadContext, TestReadWithProjectedSchemaAndSpecialFields) { std::shared_ptr read_context = std::move(unique_read_context); ASSERT_OK_AND_ASSIGN( auto internal_context, - CreateInternalReadContextForTest(read_context, table_schema, enabled_options)); + InternalReadContext::Create(read_context, table_schema, enabled_options)); auto expected_schema = DataField::ConvertDataFieldsToArrowSchema(projected_fields); ASSERT_TRUE(internal_context->GetReadSchema()->Equals(expected_schema)); } @@ -269,7 +259,7 @@ TEST(InternalReadContext, TestReadWithProjectedSchemaWithoutFieldIds) { ASSERT_OK_AND_ASSIGN( auto internal_context, - CreateInternalReadContextForTest(read_context, table_schema, table_schema->Options())); + InternalReadContext::Create(read_context, table_schema, table_schema->Options())); std::vector expected_fields = { DataField(3, arrow::field("f3", arrow::float64())), @@ -301,7 +291,7 @@ TEST(InternalReadContext, TestProjectedSchemaMetadataWhitelist) { ASSERT_OK_AND_ASSIGN( auto internal_context, - CreateInternalReadContextForTest(read_context, table_schema, table_schema->Options())); + InternalReadContext::Create(read_context, table_schema, table_schema->Options())); auto aligned_field = internal_context->GetReadSchema()->GetFieldByName("f0"); ASSERT_TRUE(aligned_field); @@ -342,7 +332,7 @@ TEST(InternalReadContext, TestMapSharedShreddingAccessRequiresSharedShreddingLay std::shared_ptr read_context = std::move(unique_read_context); ASSERT_NOK_WITH_MSG( - CreateInternalReadContextForTest(read_context, table_schema, table_schema->Options()), + InternalReadContext::Create(read_context, table_schema, table_schema->Options()), "Selected-key MAP pushdown only supports top-level shared-shredding MAP field: tags"); } diff --git a/src/paimon/format/orc/complex_predicate_test.cpp b/src/paimon/format/orc/complex_predicate_test.cpp index e7f8e1fc6..5e63e3d97 100644 --- a/src/paimon/format/orc/complex_predicate_test.cpp +++ b/src/paimon/format/orc/complex_predicate_test.cpp @@ -62,8 +62,6 @@ namespace paimon::orc::test { class ComplexPredicateTest : public ::testing::Test { public: void SetUp() override { - pool_ = GetDefaultPool(); - read_memory_ = std::make_shared(pool_); batch_size_ = 10; } void TearDown() override {} @@ -78,9 +76,11 @@ class ComplexPredicateTest : public ::testing::Test { EXPECT_OK_AND_ASSIGN(auto in_stream, OrcInputStreamImpl::Create(input_stream, DEFAULT_NATURAL_READ_SIZE)); EXPECT_TRUE(in_stream); + std::shared_ptr read_memory = + std::make_shared(GetDefaultPool()); EXPECT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), read_memory_, + OrcFileBatchReader::Create(std::move(in_stream), read_memory, /*options=*/{{"orc.timestamp-ltz.legacy.type", "false"}}, batch_size)); EXPECT_TRUE(orc_batch_reader); @@ -111,8 +111,6 @@ class ComplexPredicateTest : public ::testing::Test { } private: - std::shared_ptr pool_; - std::shared_ptr read_memory_; int32_t batch_size_; }; diff --git a/src/paimon/format/orc/orc_file_batch_reader.cpp b/src/paimon/format/orc/orc_file_batch_reader.cpp index 5be49f15b..d45cf4d9c 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader.cpp @@ -33,7 +33,6 @@ #include "orc/OrcFile.hh" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/utils/arrow/arrow_utils.h" -#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/scope_guard.h" @@ -59,10 +58,8 @@ void CollectAllColumnIds(const ::orc::Type* type, std::vector* column_ OrcFileBatchReader::OrcFileBatchReader(std::unique_ptr<::orc::ReaderMetrics>&& reader_metrics, std::unique_ptr&& reader, - const std::map& options, - const std::shared_ptr& read_memory) + const std::map& options) : options_(options), - read_memory_(read_memory), reader_metrics_(std::move(reader_metrics)), reader_(std::move(reader)), metrics_(std::make_shared()) {} @@ -96,12 +93,11 @@ Result> OrcFileBatchReader::Create( std::unique_ptr<::orc::Reader> reader = ::orc::createReader(std::move(input_stream), reader_options); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr reader_wrapper, - OrcReaderWrapper::Create(std::move(reader), file_name, batch_size, natural_read_size, - options, read_memory->arrow_pool, read_memory->orc_pool)); - auto orc_file_batch_reader = std::unique_ptr(new OrcFileBatchReader( - std::move(reader_metrics), std::move(reader_wrapper), options, read_memory)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader_wrapper, + OrcReaderWrapper::Create(std::move(reader), file_name, batch_size, + natural_read_size, options, read_memory)); + auto orc_file_batch_reader = std::unique_ptr( + new OrcFileBatchReader(std::move(reader_metrics), std::move(reader_wrapper), options)); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> file_schema, orc_file_batch_reader->GetFileSchema()); PAIMON_RETURN_NOT_OK(orc_file_batch_reader->SetReadSchema( @@ -193,7 +189,6 @@ Result OrcFileBatchReader::NextBatch() { if (BatchReader::IsEofBatch(batch)) { previous_batch_row_count_ = 0; } else { - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(batch.first.get(), read_memory_)); previous_batch_row_count_ = batch.first->length; } return batch; diff --git a/src/paimon/format/orc/orc_file_batch_reader.h b/src/paimon/format/orc/orc_file_batch_reader.h index 08f5ef27b..43810b4d1 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.h +++ b/src/paimon/format/orc/orc_file_batch_reader.h @@ -26,14 +26,11 @@ #include #include "arrow/c/bridge.h" -#include "arrow/memory_pool.h" #include "arrow/type.h" #include "orc/OrcFile.hh" #include "orc/Reader.hh" -#include "paimon/common/utils/arrow/mem_utils.h" -#include "paimon/format/orc/orc_memory_pool.h" +#include "paimon/format/orc/orc_format_defs.h" #include "paimon/format/orc/orc_reader_wrapper.h" -#include "paimon/memory/memory_pool.h" #include "paimon/predicate/predicate.h" #include "paimon/reader/prefetch_file_batch_reader.h" @@ -43,14 +40,6 @@ class InputStream; namespace paimon::orc { -struct OrcReadMemory { - explicit OrcReadMemory(const std::shared_ptr& pool) - : arrow_pool(GetSharedArrowPool(pool)), orc_pool(std::make_shared(pool)) {} - - std::shared_ptr arrow_pool; - std::shared_ptr<::orc::MemoryPool> orc_pool; -}; - class OrcFileBatchReader : public PrefetchFileBatchReader { public: ~OrcFileBatchReader() override = default; @@ -122,8 +111,7 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { private: OrcFileBatchReader(std::unique_ptr<::orc::ReaderMetrics>&& reader_metrics, std::unique_ptr&& reader, - const std::map& options, - const std::shared_ptr& read_memory); + const std::map& options); static Result<::orc::RowReaderOptions> CreateRowReaderOptions( const ::orc::Type* src_type, const ::orc::Type* target_type, @@ -136,8 +124,6 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { std::vector* target_column_ids); std::map options_; - std::shared_ptr read_memory_; - std::unique_ptr<::orc::ReaderMetrics> reader_metrics_; std::unique_ptr reader_; std::shared_ptr metrics_; diff --git a/src/paimon/format/orc/orc_file_batch_reader_test.cpp b/src/paimon/format/orc/orc_file_batch_reader_test.cpp index 79466fabe..4d303131e 100644 --- a/src/paimon/format/orc/orc_file_batch_reader_test.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader_test.cpp @@ -64,7 +64,6 @@ class OrcFileBatchReaderTest : public ::testing::Test, public: void SetUp() override { pool_ = GetDefaultPool(); - read_memory_ = std::make_shared(pool_); batch_size_ = 10; arrow::FieldVector fields = { @@ -163,9 +162,10 @@ class OrcFileBatchReaderTest : public ::testing::Test, const std::shared_ptr& predicate, const std::optional& selection_bitmap, int32_t batch_size, const std::shared_ptr& read_memory = nullptr) const { + std::shared_ptr actual_read_memory = + read_memory ? read_memory : std::make_shared(pool_); EXPECT_OK_AND_ASSIGN(auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), - read_memory ? read_memory : read_memory_, + OrcFileBatchReader::Create(std::move(in_stream), actual_read_memory, options, batch_size)); EXPECT_TRUE(orc_batch_reader); std::unique_ptr c_schema = std::make_unique(); @@ -198,7 +198,6 @@ class OrcFileBatchReaderTest : public ::testing::Test, private: std::shared_ptr pool_; - std::shared_ptr read_memory_; int32_t batch_size_; std::shared_ptr struct_array_; }; @@ -262,7 +261,8 @@ TEST_F(OrcFileBatchReaderTest, TestSetReadSchema) { std::map options = {{ORC_READ_ENABLE_LAZY_DECODING, "true"}}; ASSERT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), read_memory_, options, batch_size_)); + OrcFileBatchReader::Create(std::move(in_stream), std::make_shared(pool_), + options, batch_size_)); // test GetFileSchema() ASSERT_OK_AND_ASSIGN(auto c_file_schema, orc_batch_reader->GetFileSchema()); auto arrow_file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); @@ -638,6 +638,8 @@ TEST_F(OrcFileBatchReaderTest, TestBatchRetainsOrcReadMemory) { std::shared_ptr read_memory = std::make_shared(pool_); std::weak_ptr weak_read_memory = read_memory; + std::weak_ptr weak_arrow_pool = read_memory->arrow_pool; + std::weak_ptr<::orc::MemoryPool> weak_orc_pool = read_memory->orc_pool; std::map options = {{ORC_READ_ENABLE_LAZY_DECODING, "true"}}; std::unique_ptr reader = PrepareOrcFileBatchReader( std::move(in_stream), options, schema.get(), /*predicate=*/nullptr, @@ -647,12 +649,16 @@ TEST_F(OrcFileBatchReaderTest, TestBatchRetainsOrcReadMemory) { reader.reset(); read_memory.reset(); ASSERT_FALSE(weak_read_memory.expired()); + ASSERT_FALSE(weak_arrow_pool.expired()); + ASSERT_FALSE(weak_orc_pool.expired()); arrow::Result> import_result = arrow::ImportArray(batch.first.get(), batch.second.get()); ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); std::shared_ptr array = std::move(import_result).ValueOrDie(); ASSERT_FALSE(weak_read_memory.expired()); + ASSERT_FALSE(weak_arrow_pool.expired()); + ASSERT_FALSE(weak_orc_pool.expired()); ASSERT_TRUE(array->ValidateFull().ok()); ASSERT_OK_AND_ASSIGN( std::shared_ptr converted, @@ -661,6 +667,8 @@ TEST_F(OrcFileBatchReaderTest, TestBatchRetainsOrcReadMemory) { ASSERT_TRUE(converted->Equals(expected)); converted.reset(); ASSERT_TRUE(weak_read_memory.expired()); + ASSERT_TRUE(weak_arrow_pool.expired()); + ASSERT_TRUE(weak_orc_pool.expired()); } TEST_P(OrcFileBatchReaderTest, TestNextBatchWithTargetSchema) { @@ -704,7 +712,8 @@ TEST_F(OrcFileBatchReaderTest, TestNextBatchWithOutofOrderTargetSchema) { std::map options = {{ORC_READ_ENABLE_LAZY_DECODING, "true"}}; ASSERT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), read_memory_, options, batch_size_)); + OrcFileBatchReader::Create(std::move(in_stream), std::make_shared(pool_), + options, batch_size_)); std::unique_ptr c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(read_schema, c_schema.get()).ok()); ASSERT_NOK_WITH_MSG(orc_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, @@ -819,7 +828,8 @@ TEST_F(OrcFileBatchReaderTest, TestGetFileSchemaWithFieldId) { std::map options = {{ORC_READ_ENABLE_LAZY_DECODING, "true"}}; EXPECT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), read_memory_, options, batch_size_)); + OrcFileBatchReader::Create(std::move(in_stream), std::make_shared(pool_), + options, batch_size_)); EXPECT_TRUE(orc_batch_reader); auto c_file_schema = orc_batch_reader->GetFileSchema(); EXPECT_TRUE(c_file_schema.ok()); @@ -1284,7 +1294,8 @@ TEST_F(OrcFileBatchReaderTest, TestListStructPartialProjection) { OrcInputStreamImpl::Create(input_stream, DEFAULT_NATURAL_READ_SIZE)); ASSERT_OK_AND_ASSIGN( auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), read_memory_, /*options=*/{}, + OrcFileBatchReader::Create(std::move(in_stream), std::make_shared(pool_), + /*options=*/{}, /*batch_size=*/10)); std::unique_ptr c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(read_schema, c_schema.get()).ok()); diff --git a/src/paimon/format/orc/orc_format_defs.h b/src/paimon/format/orc/orc_format_defs.h index 4a03dc1d6..46eb5e7d8 100644 --- a/src/paimon/format/orc/orc_format_defs.h +++ b/src/paimon/format/orc/orc_format_defs.h @@ -20,8 +20,24 @@ #include #include +#include + +#include "arrow/memory_pool.h" +#include "orc/MemoryPool.hh" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/format/orc/orc_memory_pool.h" +#include "paimon/memory/memory_pool.h" namespace paimon::orc { + +struct OrcReadMemory { + explicit OrcReadMemory(const std::shared_ptr& pool) + : arrow_pool(GetSharedArrowPool(pool)), orc_pool(std::make_shared(pool)) {} + + std::shared_ptr arrow_pool; + std::shared_ptr<::orc::MemoryPool> orc_pool; +}; + // write options static inline const char ORC_STRIPE_SIZE[] = "orc.stripe.size"; static constexpr size_t DEFAULT_STRIPE_SIZE = 64 * 1024 * 1024; diff --git a/src/paimon/format/orc/orc_reader_wrapper.cpp b/src/paimon/format/orc/orc_reader_wrapper.cpp index f14826905..4eed0f500 100644 --- a/src/paimon/format/orc/orc_reader_wrapper.cpp +++ b/src/paimon/format/orc/orc_reader_wrapper.cpp @@ -25,6 +25,7 @@ #include "fmt/format.h" #include "orc/OrcFile.hh" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" @@ -81,8 +82,9 @@ Result OrcReaderWrapper::Next() { assert(orc_batch->numElements > 0); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr array, - OrcAdapter::AppendBatch(target_type_, orc_batch.get(), arrow_pool_.get())); + OrcAdapter::AppendBatch(target_type_, orc_batch.get(), read_memory_->arrow_pool.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), read_memory_)); next_row_ = GetRowNumber() + orc_batch->numElements; guard.Release(); } catch (const std::exception& e) { diff --git a/src/paimon/format/orc/orc_reader_wrapper.h b/src/paimon/format/orc/orc_reader_wrapper.h index f01ccd5a2..cc52fde19 100644 --- a/src/paimon/format/orc/orc_reader_wrapper.h +++ b/src/paimon/format/orc/orc_reader_wrapper.h @@ -34,8 +34,8 @@ #include "arrow/memory_pool.h" #include "fmt/format.h" #include "paimon/format/orc/orc_adapter.h" +#include "paimon/format/orc/orc_format_defs.h" #include "paimon/format/orc/read_range_generator.h" -#include "paimon/memory/memory_pool.h" #include "paimon/reader/batch_reader.h" namespace paimon::orc { @@ -51,14 +51,15 @@ class OrcReaderWrapper { static Result> Create( std::unique_ptr<::orc::Reader> reader, const std::string& file_name, int32_t batch_size, uint64_t natural_read_size, const std::map& options, - const std::shared_ptr& arrow_pool, - const std::shared_ptr<::orc::MemoryPool>& orc_pool) { + const std::shared_ptr& read_memory) { + if (read_memory == nullptr) { + return Status::Invalid("read memory is nullptr"); + } PAIMON_ASSIGN_OR_RAISE( std::unique_ptr range_generator, ReadRangeGenerator::Create(reader.get(), natural_read_size, options)); - return std::unique_ptr( - new OrcReaderWrapper(std::move(reader), std::move(range_generator), file_name, - batch_size, arrow_pool, orc_pool)); + return std::unique_ptr(new OrcReaderWrapper( + std::move(reader), std::move(range_generator), file_name, batch_size, read_memory)); } Status SeekToRow(uint64_t row_number); @@ -125,14 +126,12 @@ class OrcReaderWrapper { OrcReaderWrapper(std::unique_ptr<::orc::Reader> reader, std::unique_ptr range_generator, const std::string& file_name, int32_t batch_size, - const std::shared_ptr& arrow_pool, - const std::shared_ptr<::orc::MemoryPool>& orc_pool) + const std::shared_ptr& read_memory) : reader_(std::move(reader)), range_generator_(std::move(range_generator)), file_name_(file_name), batch_size_(batch_size), - arrow_pool_(arrow_pool), - orc_pool_(orc_pool) {} + read_memory_(read_memory) {} std::unique_ptr<::orc::Reader> reader_; std::unique_ptr<::orc::RowReader> row_reader_; @@ -142,8 +141,7 @@ class OrcReaderWrapper { const std::string file_name_; const int32_t batch_size_; - std::shared_ptr arrow_pool_; - std::shared_ptr<::orc::MemoryPool> orc_pool_; + std::shared_ptr read_memory_; std::shared_ptr target_type_; diff --git a/src/paimon/format/orc/orc_reader_wrapper_test.cpp b/src/paimon/format/orc/orc_reader_wrapper_test.cpp index c31ec5e45..827436010 100644 --- a/src/paimon/format/orc/orc_reader_wrapper_test.cpp +++ b/src/paimon/format/orc/orc_reader_wrapper_test.cpp @@ -64,7 +64,12 @@ TEST_F(OrcReaderWrapperTest, NextRowToRead) { writer->close(); } + std::shared_ptr read_memory = std::make_shared(GetDefaultPool()); + std::weak_ptr weak_read_memory = read_memory; + std::weak_ptr weak_arrow_pool = read_memory->arrow_pool; + std::weak_ptr<::orc::MemoryPool> weak_orc_pool = read_memory->orc_pool; ::orc::ReaderOptions reader_opts; + reader_opts.setMemoryPool(*read_memory->orc_pool); std::unique_ptr<::orc::Reader> reader = ::orc::createReader(::orc::readLocalFile(file_path), reader_opts); std::map options; @@ -74,8 +79,8 @@ TEST_F(OrcReaderWrapperTest, NextRowToRead) { /*batch_size=*/2, /*natural_read_size=*/0, /*options=*/options, - /*arrow_pool=*/GetArrowPool(GetDefaultPool()), - /*orc_pool=*/nullptr)); + /*read_memory=*/read_memory)); + read_memory.reset(); auto data_types = arrow::struct_({arrow::field("col1", arrow::int64()), arrow::field("col2", arrow::utf8())}); ::orc::RowReaderOptions row_opts; @@ -87,11 +92,19 @@ TEST_F(OrcReaderWrapperTest, NextRowToRead) { ASSERT_OK_AND_ASSIGN(auto batch2, wrapper->Next()); EXPECT_EQ(wrapper->GetNextRowToRead(), 3u); // only 1 row left - ReaderUtils::ReleaseReadBatch(std::move(batch2)); ASSERT_OK_AND_ASSIGN(auto batch3, wrapper->Next()); EXPECT_EQ(wrapper->GetNextRowToRead(), 3u); ReaderUtils::ReleaseReadBatch(std::move(batch3)); + + wrapper.reset(); + EXPECT_FALSE(weak_read_memory.expired()); + EXPECT_FALSE(weak_arrow_pool.expired()); + EXPECT_FALSE(weak_orc_pool.expired()); + ReaderUtils::ReleaseReadBatch(std::move(batch2)); + EXPECT_TRUE(weak_read_memory.expired()); + EXPECT_TRUE(weak_arrow_pool.expired()); + EXPECT_TRUE(weak_orc_pool.expired()); } } // namespace paimon::orc::test diff --git a/src/paimon/format/orc/predicate_pushdown_test.cpp b/src/paimon/format/orc/predicate_pushdown_test.cpp index 100c05b79..b19cabe6f 100644 --- a/src/paimon/format/orc/predicate_pushdown_test.cpp +++ b/src/paimon/format/orc/predicate_pushdown_test.cpp @@ -69,7 +69,6 @@ class PredicatePushdownTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); - read_memory_ = std::make_shared(pool_); batch_size_ = 10; arrow::FieldVector fields = { @@ -114,8 +113,9 @@ class PredicatePushdownTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_name_)); ASSERT_OK_AND_ASSIGN(auto in_stream, OrcInputStreamImpl::Create(in, DEFAULT_NATURAL_READ_SIZE)); + std::shared_ptr read_memory = std::make_shared(pool_); ASSERT_OK_AND_ASSIGN(auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), read_memory_, + OrcFileBatchReader::Create(std::move(in_stream), read_memory, /*options=*/{}, batch_size_)); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); @@ -141,7 +141,6 @@ class PredicatePushdownTest : public ::testing::Test { private: std::shared_ptr pool_; - std::shared_ptr read_memory_; int32_t batch_size_; std::shared_ptr struct_array_; std::shared_ptr fs_; @@ -421,8 +420,9 @@ TEST_F(PredicatePushdownTest, TestPredicatePushdownWithNullLiteral) { ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_name_)); ASSERT_OK_AND_ASSIGN(auto in_stream, OrcInputStreamImpl::Create(in, DEFAULT_NATURAL_READ_SIZE)); + std::shared_ptr read_memory = std::make_shared(pool_); ASSERT_OK_AND_ASSIGN(auto orc_batch_reader, - OrcFileBatchReader::Create(std::move(in_stream), read_memory_, + OrcFileBatchReader::Create(std::move(in_stream), read_memory, /*options=*/{}, batch_size_)); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); From 1332736042e26ec4d90011387f34ba815e81b243 Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Sat, 29 Aug 2026 23:27:07 +0800 Subject: [PATCH 4/8] fix sort merge --- .../core/io/concat_key_value_record_reader.h | 23 +++-------- .../concat_key_value_record_reader_test.cpp | 34 +--------------- .../core/mergetree/compact/loser_tree.cpp | 19 ++------- .../core/mergetree/compact/loser_tree.h | 23 ++++++----- .../compact/sort_merge_reader_test.cpp | 39 ++----------------- .../sort_merge_reader_with_min_heap.cpp | 25 +++++------- .../compact/sort_merge_reader_with_min_heap.h | 29 +++++++------- 7 files changed, 46 insertions(+), 146 deletions(-) diff --git a/src/paimon/core/io/concat_key_value_record_reader.h b/src/paimon/core/io/concat_key_value_record_reader.h index 7a3f545ed..99318a450 100644 --- a/src/paimon/core/io/concat_key_value_record_reader.h +++ b/src/paimon/core/io/concat_key_value_record_reader.h @@ -37,7 +37,7 @@ class ConcatKeyValueRecordReader : public KeyValueRecordReader { public: explicit ConcatKeyValueRecordReader( std::vector>&& readers) - : finished_reader_metrics_(std::make_shared()), readers_(std::move(readers)) {} + : readers_(std::move(readers)) {} Result> NextBatch() override { while (current_ < readers_.size()) { @@ -48,7 +48,7 @@ class ConcatKeyValueRecordReader : public KeyValueRecordReader { return iterator; } // current meets eof, move to next reader - CloseAndReleaseReader(current_); + current_reader->Close(); current_++; } // read finish @@ -57,29 +57,16 @@ class ConcatKeyValueRecordReader : public KeyValueRecordReader { void Close() override { for (; current_ < readers_.size(); current_++) { - CloseAndReleaseReader(current_); + readers_[current_]->Close(); } } std::shared_ptr GetReaderMetrics() const override { - auto metrics = std::make_shared(); - metrics->Merge(finished_reader_metrics_); - metrics->Merge(MetricsImpl::CollectReadMetrics(readers_)); - return metrics; + return MetricsImpl::CollectReadMetrics(readers_); } private: - void CloseAndReleaseReader(size_t reader_index) { - std::unique_ptr& reader = readers_[reader_index]; - if (!reader) { - return; - } - reader->Close(); - finished_reader_metrics_->Merge(reader->GetReaderMetrics()); - reader.reset(); - } - - std::shared_ptr finished_reader_metrics_; + // KeyValue rows may outlive the active child and still reference buffers allocated by it. std::vector> readers_; size_t current_{0}; }; diff --git a/src/paimon/core/io/concat_key_value_record_reader_test.cpp b/src/paimon/core/io/concat_key_value_record_reader_test.cpp index b9a2ac234..b9da3cfd3 100644 --- a/src/paimon/core/io/concat_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/concat_key_value_record_reader_test.cpp @@ -37,30 +37,6 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { - -class ConcatLifetimeTrackingKeyValueRecordReader : public KeyValueRecordReader { - public: - ConcatLifetimeTrackingKeyValueRecordReader(std::unique_ptr reader, - const std::shared_ptr& lifetime) - : reader_(std::move(reader)), lifetime_(lifetime) {} - - Result> NextBatch() override { - return reader_->NextBatch(); - } - - std::shared_ptr GetReaderMetrics() const override { - return reader_->GetReaderMetrics(); - } - - void Close() override { - reader_->Close(); - } - - private: - std::unique_ptr reader_; - std::shared_ptr lifetime_; -}; - class ConcatKeyValueRecordReaderTest : public testing::Test { public: void SetUp() override { @@ -73,18 +49,13 @@ class ConcatKeyValueRecordReaderTest : public testing::Test { const std::shared_ptr& value_schema) const { for (auto batch_size : {1, 2, 3, 4, 100}) { std::vector> record_readers; - std::vector> reader_lifetimes; for (const auto& src_array : src_array_vec) { auto src_type = src_array->type(); auto file_batch_reader = std::make_unique( src_array, src_type, /*batch_size=*/batch_size); auto record_reader = std::make_unique( std::move(file_batch_reader), key_schema, value_schema, /*level=*/0, pool_); - std::shared_ptr lifetime = std::make_shared(0); - reader_lifetimes.push_back(lifetime); - record_readers.push_back( - std::make_unique( - std::move(record_reader), lifetime)); + record_readers.push_back(std::move(record_reader)); } auto concat_record_reader = @@ -96,9 +67,6 @@ class ConcatKeyValueRecordReaderTest : public testing::Test { concat_record_reader.get()))); KeyValueChecker::CheckResult(expected, results, key_schema->num_fields(), value_schema->num_fields()); - for (const auto& reader_lifetime : reader_lifetimes) { - ASSERT_TRUE(reader_lifetime.expired()); - } } } diff --git a/src/paimon/core/mergetree/compact/loser_tree.cpp b/src/paimon/core/mergetree/compact/loser_tree.cpp index 596cbc5b8..e93211903 100644 --- a/src/paimon/core/mergetree/compact/loser_tree.cpp +++ b/src/paimon/core/mergetree/compact/loser_tree.cpp @@ -21,32 +21,19 @@ #include #include -#include "paimon/common/metrics/metrics_impl.h" - namespace paimon { LoserTree::LoserTree(std::vector>&& readers, const CompareFunc& first_comparator, const CompareFunc& second_comparator) : size_(readers.size()), initialized_(false), + readers_holder_(std::move(readers)), tree_(size_), first_comparator_(first_comparator), second_comparator_(second_comparator) { leaves_.reserve(size_); - for (auto& reader : readers) { - leaves_.emplace_back(std::move(reader)); - } -} - -std::shared_ptr LoserTree::GetReaderMetrics() const { - auto metrics = std::make_shared(); - for (const auto& leaf : leaves_) { - if (leaf.reader) { - metrics->Merge(leaf.reader->GetReaderMetrics()); - } else { - metrics->Merge(leaf.finished_reader_metrics); - } + for (const auto& reader : readers_holder_) { + leaves_.emplace_back(reader.get()); } - return metrics; } Status LoserTree::InitializeIfNeeded() { diff --git a/src/paimon/core/mergetree/compact/loser_tree.h b/src/paimon/core/mergetree/compact/loser_tree.h index 1bde92131..50b3bab5a 100644 --- a/src/paimon/core/mergetree/compact/loser_tree.h +++ b/src/paimon/core/mergetree/compact/loser_tree.h @@ -24,6 +24,7 @@ #include #include +#include "paimon/common/metrics/metrics_impl.h" #include "paimon/core/io/key_value_record_reader.h" #include "paimon/core/key_value.h" #include "paimon/result.h" @@ -51,13 +52,13 @@ class LoserTree { /// Peek the current winner, mainly for key comparisons. const std::optional& PeekWinner() const; - std::shared_ptr GetReaderMetrics() const; + std::shared_ptr GetReaderMetrics() const { + return MetricsImpl::CollectReadMetrics(readers_holder_); + } void Close() { - for (const auto& leaf : leaves_) { - if (leaf.reader) { - leaf.reader->Close(); - } + for (const auto& reader : readers_holder_) { + reader->Close(); } } @@ -96,8 +97,7 @@ class LoserTree { } struct LeafIterator { - explicit LeafIterator(std::unique_ptr&& reader) - : reader(std::move(reader)) {} + explicit LeafIterator(KeyValueRecordReader* reader) : reader(reader) {} const std::optional& Peek() const { return kv; @@ -127,10 +127,7 @@ class LoserTree { PAIMON_ASSIGN_OR_RAISE(iterator, reader->NextBatch()); if (!iterator) { // read eof - iterator.reset(); reader->Close(); - finished_reader_metrics = reader->GetReaderMetrics(); - reader.reset(); end_of_input = true; kv = std::nullopt; } else { @@ -151,8 +148,7 @@ class LoserTree { bool end_of_input = false; int32_t first_same_key_index = -1; State state = State::WINNER_WITH_NEW_KEY; - std::unique_ptr reader; - std::shared_ptr finished_reader_metrics; + KeyValueRecordReader* reader; std::unique_ptr iterator; std::optional kv; }; @@ -160,6 +156,9 @@ class LoserTree { private: int32_t size_; bool initialized_; + // KeyValue rows may be consumed asynchronously and still reference buffers allocated by the + // input readers. + std::vector> readers_holder_; std::vector tree_; std::vector leaves_; diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp b/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp index cc7bf5e15..dabee55c4 100644 --- a/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp +++ b/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp @@ -52,30 +52,6 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { - -class SortMergeLifetimeTrackingKeyValueRecordReader : public KeyValueRecordReader { - public: - SortMergeLifetimeTrackingKeyValueRecordReader(std::unique_ptr reader, - const std::shared_ptr& lifetime) - : reader_(std::move(reader)), lifetime_(lifetime) {} - - Result> NextBatch() override { - return reader_->NextBatch(); - } - - std::shared_ptr GetReaderMetrics() const override { - return reader_->GetReaderMetrics(); - } - - void Close() override { - reader_->Close(); - } - - private: - std::unique_ptr reader_; - std::shared_ptr lifetime_; -}; - class SortMergeReaderTest : public testing::Test { public: void SetUp() override { @@ -116,7 +92,7 @@ class SortMergeReaderTest : public testing::Test { const std::shared_ptr& user_defined_seq_comparator, const std::shared_ptr& key_schema, const std::shared_ptr& value_schema, int32_t batch_size, bool need_merge, - bool ignore_delete, std::vector>* reader_lifetimes) const { + bool ignore_delete) const { auto mfunc = std::make_unique(ignore_delete); auto merge_function_wrapper = std::make_shared(std::move(mfunc)); @@ -135,13 +111,8 @@ class SortMergeReaderTest : public testing::Test { std::move(file_batch_reader), key_schema, value_schema, /*level=*/0, pool_); std::vector> readers; readers.push_back(std::move(record_reader)); - std::unique_ptr concat_reader = - std::make_unique(std::move(readers)); - std::shared_ptr lifetime = std::make_shared(0); - reader_lifetimes->push_back(lifetime); concat_readers.push_back( - std::make_unique( - std::move(concat_reader), lifetime)); + std::make_unique(std::move(readers))); } return std::make_unique(std::move(concat_readers), user_key_comparator, @@ -158,19 +129,15 @@ class SortMergeReaderTest : public testing::Test { const std::vector& expected, bool need_merge, bool ignore_delete = false) const { for (auto batch_size : {1, 2, 3, 4, 100}) { - std::vector> reader_lifetimes; auto sort_merge_reader = CreateSortMergeReader( src_array_vec, user_key_comparator, user_defined_seq_comparator, key_schema, - value_schema, batch_size, need_merge, ignore_delete, &reader_lifetimes); + value_schema, batch_size, need_merge, ignore_delete); ASSERT_OK_AND_ASSIGN( std::vector results, (ReadResultCollector::CollectKeyValueResult< SortMergeReader, SortMergeReader::Iterator>(sort_merge_reader.get()))); KeyValueChecker::CheckResult(expected, results, key_schema->num_fields(), value_schema->num_fields()); - for (const auto& reader_lifetime : reader_lifetimes) { - ASSERT_TRUE(reader_lifetime.expired()); - } } } diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp index fc592accf..08629e638 100644 --- a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp +++ b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp @@ -30,40 +30,35 @@ SortMergeReaderWithMinHeap::SortMergeReaderWithMinHeap( const std::shared_ptr& user_defined_seq_comparator, const std::shared_ptr>& merge_function_wrapper) : need_merge_(merge_function_wrapper != nullptr), - readers_(std::move(readers)), - finished_reader_metrics_(std::make_shared()), + readers_holder_(std::move(readers)), user_key_comparator_(user_key_comparator), merge_function_wrapper_(merge_function_wrapper), min_heap_(HeapSorter(user_key_comparator, user_defined_seq_comparator)) { - next_batch_reader_indices_.reserve(readers_.size()); - for (size_t reader_index = 0; reader_index < readers_.size(); reader_index++) { - next_batch_reader_indices_.push_back(reader_index); + next_batch_readers_.reserve(readers_holder_.size()); + for (auto& reader : readers_holder_) { + next_batch_readers_.push_back(reader.get()); } } Result> SortMergeReaderWithMinHeap::NextBatch() { - for (size_t reader_index : next_batch_reader_indices_) { - std::unique_ptr& reader = readers_[reader_index]; - assert(reader); + for (auto* reader : next_batch_readers_) { while (true) { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, reader->NextBatch()); if (!iterator) { // no more batches, permanently remove this reader reader->Close(); - finished_reader_metrics_->Merge(reader->GetReaderMetrics()); - reader.reset(); break; } PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext()); if (has_next) { PAIMON_ASSIGN_OR_RAISE(KeyValue kv, iterator->Next()); - min_heap_.emplace(std::move(kv), std::move(iterator), reader_index); + min_heap_.emplace(std::move(kv), std::move(iterator), reader); break; } } } - next_batch_reader_indices_.clear(); + next_batch_readers_.clear(); if (min_heap_.empty()) { return std::unique_ptr(); } @@ -90,7 +85,7 @@ Result SortMergeReaderWithMinHeap::Iterator::HasNext() { } Result SortMergeReaderWithMinHeap::Iterator::NextImpl() { - assert(reader_->next_batch_reader_indices_.empty()); + assert(reader_->next_batch_readers_.empty()); // add previously polled elements back to priority queue for (auto& element : reader_->polled_) { PAIMON_ASSIGN_OR_RAISE(bool updated, element.Update()); @@ -99,11 +94,11 @@ Result SortMergeReaderWithMinHeap::Iterator::NextImpl() { reader_->min_heap_.push(std::move(element)); } else { // reach end of batch, clean up - reader_->next_batch_reader_indices_.push_back(element.reader_index); + reader_->next_batch_readers_.push_back(std::move(element.reader)); } } reader_->polled_.clear(); - if (!reader_->next_batch_reader_indices_.empty()) { + if (!reader_->next_batch_readers_.empty()) { return false; } diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h index 667af1242..4bcf8fb2a 100644 --- a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h +++ b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h @@ -69,10 +69,8 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { Result> NextBatch() override; void Close() override { - for (const auto& reader : readers_) { - if (reader) { - reader->Close(); - } + for (const auto& reader : readers_holder_) { + reader->Close(); } if (merge_function_wrapper_) { merge_function_wrapper_->Reset(); @@ -80,23 +78,21 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { } std::shared_ptr GetReaderMetrics() const override { - auto metrics = std::make_shared(); - metrics->Merge(finished_reader_metrics_); - metrics->Merge(MetricsImpl::CollectReadMetrics(readers_)); - return metrics; + return MetricsImpl::CollectReadMetrics(readers_holder_); } private: struct Element { Element(KeyValue&& _kv, std::unique_ptr&& _iterator, - size_t _reader_index) - : kv(std::move(_kv)), reader_index(_reader_index), iterator(std::move(_iterator)) { + KeyValueRecordReader* _reader) + : kv(std::move(_kv)), reader(_reader), iterator(std::move(_iterator)) { assert(iterator); + assert(reader); } Element(Element&& other) noexcept : kv(std::move(other.kv)) { iterator = std::move(other.iterator); - reader_index = other.reader_index; + reader = other.reader; } Element& operator=(Element&& other) noexcept { @@ -105,7 +101,7 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { } kv = std::move(other.kv); iterator = std::move(other.iterator); - reader_index = other.reader_index; + reader = other.reader; return *this; } @@ -121,7 +117,7 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { public: KeyValue kv; - size_t reader_index; + KeyValueRecordReader* reader; std::unique_ptr iterator; }; @@ -153,9 +149,10 @@ class SortMergeReaderWithMinHeap : public SortMergeReader { private: const bool need_merge_; - std::vector> readers_; - std::vector next_batch_reader_indices_; - std::shared_ptr finished_reader_metrics_; + // KeyValue rows may be consumed asynchronously and still reference buffers allocated by the + // input readers. + std::vector> readers_holder_; + std::vector next_batch_readers_; std::shared_ptr user_key_comparator_; std::shared_ptr> merge_function_wrapper_; std::priority_queue, HeapSorter> min_heap_; From bdbcdf4098b889fb62934f10f36af3212149574d Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Mon, 31 Aug 2026 08:52:24 +0800 Subject: [PATCH 5/8] fix clang-tidy --- src/paimon/common/utils/arrow/mem_utils.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/paimon/common/utils/arrow/mem_utils.cpp b/src/paimon/common/utils/arrow/mem_utils.cpp index f454bda25..94a7339fc 100644 --- a/src/paimon/common/utils/arrow/mem_utils.cpp +++ b/src/paimon/common/utils/arrow/mem_utils.cpp @@ -142,8 +142,7 @@ Status AddArrowArrayLifetime(ArrowArray* array, const std::shared_ptr& lif return Status::Invalid("cannot add an empty lifetime to an ArrowArray"); } if (array->release == ReleaseArrowArray) { - const ArrowArrayPrivateData* data = - static_cast(array->private_data); + const auto* data = static_cast(array->private_data); if (HasSameOwner(data->lifetime, lifetime)) { return Status::OK(); } From e167dbff8ad5baf5612b26103896197ac962de25 Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Mon, 31 Aug 2026 15:31:20 +0800 Subject: [PATCH 6/8] rebase realtime --- src/paimon/common/utils/arrow/mem_utils.cpp | 3 +- src/paimon/common/utils/arrow/mem_utils.h | 3 +- .../common/utils/arrow/mem_utils_test.cpp | 5 ++- .../core/io/vector_file_batch_reader_test.cpp | 30 +++++++++------- .../realtime/arrow_realtime_store_test.cpp | 1 + .../realtime/primary_key_realtime_store.cpp | 4 +-- .../primary_key_realtime_store_test.cpp | 34 ++++++++---------- .../realtime/realtime_primary_key_writer.cpp | 4 +-- .../avro/avro_file_batch_reader_test.cpp | 1 + .../format/mosaic/mosaic_file_format_test.cpp | 1 + .../format/orc/orc_file_batch_reader_test.cpp | 1 + .../parquet_file_batch_reader_test.cpp | 1 + .../format/parquet/variant_parquet_test.cpp | 4 +-- test/inte/realtime_write_inte_test.cpp | 35 ++++++++----------- 14 files changed, 67 insertions(+), 60 deletions(-) diff --git a/src/paimon/common/utils/arrow/mem_utils.cpp b/src/paimon/common/utils/arrow/mem_utils.cpp index 94a7339fc..99eab4ce5 100644 --- a/src/paimon/common/utils/arrow/mem_utils.cpp +++ b/src/paimon/common/utils/arrow/mem_utils.cpp @@ -139,6 +139,7 @@ Status AddArrowArrayLifetime(ArrowArray* array, const std::shared_ptr& lif return Status::Invalid("cannot add lifetime to a released ArrowArray"); } if (lifetime.use_count() == 0) { + ArrowArrayRelease(array); return Status::Invalid("cannot add an empty lifetime to an ArrowArray"); } if (array->release == ReleaseArrowArray) { @@ -154,7 +155,7 @@ Status AddArrowArrayLifetime(ArrowArray* array, const std::shared_ptr& lif array->private_data = data.release(); return Status::OK(); } catch (const std::bad_alloc&) { - array->release(array); + ArrowArrayRelease(array); return Status::OutOfMemory("failed to add lifetime to an ArrowArray"); } } diff --git a/src/paimon/common/utils/arrow/mem_utils.h b/src/paimon/common/utils/arrow/mem_utils.h index 5f81437aa..101245b57 100644 --- a/src/paimon/common/utils/arrow/mem_utils.h +++ b/src/paimon/common/utils/arrow/mem_utils.h @@ -41,7 +41,8 @@ PAIMON_EXPORT std::shared_ptr GetSharedArrowPool( /// /// The existing release callback and private data are preserved as a release chain, so this helper /// can be applied to arrays produced by Arrow libraries or format plugins without inspecting their -/// private data. +/// private data. Once a valid, unreleased array is passed in, this function releases it if +/// retaining the lifetime fails. PAIMON_EXPORT Status AddArrowArrayLifetime(ArrowArray* array, const std::shared_ptr& lifetime); diff --git a/src/paimon/common/utils/arrow/mem_utils_test.cpp b/src/paimon/common/utils/arrow/mem_utils_test.cpp index 820f3b95a..530572169 100644 --- a/src/paimon/common/utils/arrow/mem_utils_test.cpp +++ b/src/paimon/common/utils/arrow/mem_utils_test.cpp @@ -283,9 +283,12 @@ TEST(MemUtilsTest, TestAddArrowArrayLifetimeRejectsInvalidInput) { ASSERT_NOK(AddArrowArrayLifetime(nullptr, std::make_shared(1))); ASSERT_NOK(AddArrowArrayLifetime(&array, std::make_shared(1))); + std::vector release_order; array.release = ReleaseTrackingArrowArray; + array.private_data = new TrackingArrowArrayPrivateData{nullptr, &release_order}; ASSERT_NOK(AddArrowArrayLifetime(&array, nullptr)); - array.release = nullptr; + ASSERT_EQ(nullptr, array.release); + ASSERT_EQ(std::vector({0}), release_order); } } // namespace paimon::test diff --git a/src/paimon/core/io/vector_file_batch_reader_test.cpp b/src/paimon/core/io/vector_file_batch_reader_test.cpp index 5eff19e73..ec86a1487 100644 --- a/src/paimon/core/io/vector_file_batch_reader_test.cpp +++ b/src/paimon/core/io/vector_file_batch_reader_test.cpp @@ -98,15 +98,17 @@ TEST(VectorFileBatchReaderTest, KeepFixedSizeListFileSchema) { std::make_unique(logical_array, logical_type, /*batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); MockFileBatchReader* inner_reader = mock_reader.get(); - VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); + auto reader = std::make_unique(std::move(mock_reader), + GetSharedArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); - ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, - /*selection_bitmap=*/std::nullopt)); + ASSERT_OK(reader->SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); ASSERT_EQ(inner_reader->read_schema_->field(1)->type()->id(), arrow::Type::FIXED_SIZE_LIST); - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader.NextBatch()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + reader.reset(); arrow::Result> actual_result = arrow::ImportArray(batch.first.get(), batch.second.get()); ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); @@ -136,14 +138,16 @@ TEST(VectorFileBatchReaderTest, NormalizeFixedSizeListElementField) { auto mock_reader = std::make_unique(file_array, file_type, /*batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); - VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); + auto reader = std::make_unique(std::move(mock_reader), + GetSharedArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); - ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, - /*selection_bitmap=*/std::nullopt)); + ASSERT_OK(reader->SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader.NextBatch()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + reader.reset(); arrow::Result> actual_result = arrow::ImportArray(batch.first.get(), batch.second.get()); ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); @@ -174,16 +178,18 @@ TEST(VectorFileBatchReaderTest, ConvertNestedVectorsWithBitmap) { auto mock_reader = std::make_unique(physical_array, physical_type, bitmap, /*read_batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); - VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); + auto reader = std::make_unique(std::move(mock_reader), + GetSharedArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); - ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, - /*selection_bitmap=*/std::nullopt)); + ASSERT_OK(reader->SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, - reader.NextBatchWithBitmap()); + reader->NextBatchWithBitmap()); ASSERT_FALSE(batch_with_bitmap.second.Contains(0)); ASSERT_TRUE(batch_with_bitmap.second.Contains(1)); + reader.reset(); arrow::Result> actual_result = arrow::ImportArray( batch_with_bitmap.first.first.get(), batch_with_bitmap.first.second.get()); ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); diff --git a/src/paimon/core/realtime/arrow_realtime_store_test.cpp b/src/paimon/core/realtime/arrow_realtime_store_test.cpp index 864b1f810..ef76c7d60 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_test.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_test.cpp @@ -213,6 +213,7 @@ TEST_F(ArrowRealtimeStoreTest, TestCommitReaderPreservesSlicedBatch) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + readers.clear(); arrow::Result> import_result = arrow::ImportArray(batch.first.get(), batch.second.get()); ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index 421eb0c8d..3e46e0788 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -112,7 +112,7 @@ class StoredBatchReader final : public BatchReader { ArrowUtils::NormalizeRecordBatchOffsets(record_batch, arrow_pool_.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportRecordBatch(*normalized_batch, array.get(), schema.get())); - PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(array.get(), arrow_pool_)); data_.reset(); arrow_pool_.reset(); export_guard.Release(); @@ -276,7 +276,7 @@ Result> PrimaryKeyRealtimeStore::Create if (!memory_pool) { return Status::Invalid("PK real-time store memory pool is null"); } - std::shared_ptr arrow_pool = GetArrowPool(memory_pool); + std::shared_ptr arrow_pool = GetSharedArrowPool(memory_pool); return std::shared_ptr(new PrimaryKeyRealtimeStore( std::make_unique(transport_schema, std::move(arrow_pool)))); } diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index ed80db275..244855037 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -35,6 +35,7 @@ #include "paimon/macros.h" #include "paimon/memory/memory_pool.h" #include "paimon/realtime/arrow_realtime_store_factory.h" +#include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -94,18 +95,13 @@ void AssertOffsetsZero(const ArrowArray* array) { } } -Result ReadJson(const std::vector>& readers) { +Result ReadJson(std::vector> readers) { std::vector> batches; - for (const std::unique_ptr& reader : readers) { - while (true) { - PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); - if (BatchReader::IsEofBatch(batch)) { - break; - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr array, - arrow::ImportArray(batch.first.get(), batch.second.get())); - batches.push_back(std::move(array)); + for (std::unique_ptr& reader : readers) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, + ReadResultCollector::CollectResult(std::move(reader))); + if (result) { + batches.insert(batches.end(), result->chunks().begin(), result->chunks().end()); } } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, @@ -151,16 +147,13 @@ TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderPerStoredBatch) { ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateCommitReaders(segment.value())); ASSERT_EQ(2, readers.size()); - ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(std::move(readers))); ASSERT_EQ( "-- is_valid: all not null\n-- child 0 type: int8\n [\n 1,\n 0,\n 2\n ]\n-- " "child 1 type: int64\n [\n 6,\n 5,\n 7\n ]\n-- child 2 type: int64\n [\n " "1,\n 0,\n 2\n ]\n-- child 3 type: int64\n [\n 1,\n 3,\n 2\n ]\n-- child " "4 type: string\n [\n \"before\",\n \"three\",\n \"after\"\n ]", actual); - for (const std::unique_ptr& reader : readers) { - reader->Close(); - } } void AssertSlicedBatch(BatchReader* reader) { @@ -264,7 +257,7 @@ TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { /*enable_predicate_pushdown=*/false}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(retained_view, /*offset_begin=*/0, context)); - ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(std::move(readers))); ASSERT_NE(std::string::npos, actual.find("\"one\"")); ASSERT_NE(std::string::npos, actual.find("\"two\"")); ASSERT_NE(std::string::npos, actual.find("\"three\"")); @@ -288,12 +281,12 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(view, /*offset_begin=*/0, context)); ASSERT_EQ(2, readers.size()); - ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(std::move(readers))); ASSERT_NE(std::string::npos, actual.find("\"one\"")); ASSERT_NE(std::string::npos, actual.find("\"two\"")); } -TEST(PrimaryKeyRealtimeStoreTest, TestQueryPoolOutlivesStoreReaderAndExport) { +TEST(PrimaryKeyRealtimeStoreTest, TestQueryBatchOutlivesStoreAndReader) { const std::shared_ptr stored_schema = TransportSchema(); std::shared_ptr pool = GetMemoryPool(); std::weak_ptr pool_lifetime = pool; @@ -322,11 +315,13 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryPoolOutlivesStoreReaderAndExport) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + readers.clear(); + ASSERT_FALSE(pool_lifetime.expired()); + arrow::Result> import_result = arrow::ImportArray(batch.first.get(), batch.second.get()); ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); std::shared_ptr imported = std::move(import_result).ValueOrDie(); - readers.clear(); ASSERT_FALSE(pool_lifetime.expired()); imported.reset(); ASSERT_TRUE(pool_lifetime.expired()); @@ -371,6 +366,7 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderProjectsNestedFields) { ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(view, /*offset_begin=*/0, context)); ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + readers.clear(); arrow::Result> import_result = arrow::ImportArray(batch.first.get(), batch.second.get()); ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index 61cd14009..b2c04c4c5 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -158,7 +158,7 @@ RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( int64_t next_offset, int64_t last_sequence_number, const std::shared_ptr& memory_pool) : memory_pool_(memory_pool), - arrow_pool_(GetArrowPool(memory_pool)), + arrow_pool_(GetSharedArrowPool(memory_pool)), realtime_store_(realtime_store), merge_tree_writer_(merge_tree_writer), realtime_context_(realtime_context), @@ -205,7 +205,7 @@ Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { arrow_pool_.get())); auto output = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*transport, output.get())); - PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(output.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(output.get(), arrow_pool_)); RecordBatchBuilder builder(output.get()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr transport_batch, builder.Finish()); PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ diff --git a/src/paimon/format/avro/avro_file_batch_reader_test.cpp b/src/paimon/format/avro/avro_file_batch_reader_test.cpp index 476af7bde..bd0d63eb6 100644 --- a/src/paimon/format/avro/avro_file_batch_reader_test.cpp +++ b/src/paimon/format/avro/avro_file_batch_reader_test.cpp @@ -411,6 +411,7 @@ TEST_F(AvroFileBatchReaderTest, TestSetReadSchemaResetsReaderToFirstRow) { ASSERT_OK_AND_ASSIGN(auto projected_batch, reader->NextBatch()); ASSERT_EQ(0, reader->GetPreviousBatchFileRowId(0).value()); + reader.reset(); auto projected_array = arrow::ImportArray(projected_batch.first.get(), projected_batch.second.get()).ValueOrDie(); auto expected_projected_array = arrow::ipc::internal::json::ArrayFromJSON( diff --git a/src/paimon/format/mosaic/mosaic_file_format_test.cpp b/src/paimon/format/mosaic/mosaic_file_format_test.cpp index d356fb26b..deb951191 100644 --- a/src/paimon/format/mosaic/mosaic_file_format_test.cpp +++ b/src/paimon/format/mosaic/mosaic_file_format_test.cpp @@ -267,6 +267,7 @@ TEST_F(MosaicFileFormatTest, SetReadSchemaResetsReaderToFirstRow) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch projected_batch, reader->NextBatch()); ASSERT_OK_AND_ASSIGN(first_row, reader->GetPreviousBatchFileRowId(/*batch_row_id=*/0)); ASSERT_EQ(first_row, 0); + reader.reset(); std::shared_ptr projected_array = arrow::ImportArray(projected_batch.first.get(), projected_batch.second.get()).ValueOrDie(); std::shared_ptr expected = diff --git a/src/paimon/format/orc/orc_file_batch_reader_test.cpp b/src/paimon/format/orc/orc_file_batch_reader_test.cpp index 4d303131e..048cdc0ff 100644 --- a/src/paimon/format/orc/orc_file_batch_reader_test.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader_test.cpp @@ -943,6 +943,7 @@ TEST_F(OrcFileBatchReaderTest, TestReadNoField) { ASSERT_NOK(orc_batch_reader->GetPreviousBatchFileRowId(0)); ASSERT_TRUE(BatchReader::IsEofBatch(batch4)); orc_batch_reader->Close(); + orc_batch_reader.reset(); arrow::FieldVector fields; auto arrow_type = arrow::struct_(fields); diff --git a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp index 86eab1458..6190a955c 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -1214,6 +1214,7 @@ TEST_F(ParquetFileBatchReaderTest, TestReadNoField) { ASSERT_TRUE(BatchReader::IsEofBatch(batch4)); ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); parquet_batch_reader->Close(); + parquet_batch_reader.reset(); arrow::FieldVector fields; auto arrow_type = arrow::struct_(fields); diff --git a/src/paimon/format/parquet/variant_parquet_test.cpp b/src/paimon/format/parquet/variant_parquet_test.cpp index e0e73df15..a1481dcfb 100644 --- a/src/paimon/format/parquet/variant_parquet_test.cpp +++ b/src/paimon/format/parquet/variant_parquet_test.cpp @@ -432,12 +432,12 @@ class VariantParquetTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN(auto batch_with_bitmap, shredding_reader->NextBatchWithBitmap()); ASSERT_FALSE(BatchReader::IsEofBatch(batch_with_bitmap)); auto& [read_batch, bitmap] = batch_with_bitmap; + shredding_reader->Close(); + shredding_reader.reset(); auto imported = arrow::ImportArray(read_batch.first.get(), read_batch.second.get()); ASSERT_TRUE(imported.ok()) << imported.status().ToString(); auto result_struct = checked_pointer_cast(imported.ValueOrDie()); *column = result_struct->field(1); - shredding_reader->Close(); - shredding_reader.reset(); } // `ReadColumn` for the cases whose second column is a struct. diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index bba491ab4..5bc7bcb6c 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -1034,14 +1034,11 @@ TEST_F(RealtimeWriteInteTest, TestPkRead) { Literal(FieldType::STRING, expected_payload.data(), expected_payload.size())); ASSERT_OK_AND_ASSIGN(std::shared_ptr filtered_plan, CreatePlan(realtime_context, predicate)); - ASSERT_OK_AND_ASSIGN( - CollectedReadResult filtered_result, - ReadPlan(filtered_plan, realtime_context, {"id", "payload", "pt"}, predicate, - /*enable_predicate_filter=*/true)); - ASSERT_EQ(nullptr, filtered_result.data); + ASSERT_OK_AND_ASSIGN(std::shared_ptr filtered_result, + ReadPlan(filtered_plan, realtime_context, {"id", "payload", "pt"}, + predicate, /*enable_predicate_filter=*/true)); + ASSERT_EQ(nullptr, filtered_result); ASSERT_FALSE(saw_query_predicate->load(std::memory_order_acquire)); - filtered_result.reader->Close(); - filtered_result.reader.reset(); ASSERT_OK(writer->Close()); writer.reset(); @@ -1067,13 +1064,13 @@ TEST_F(RealtimeWriteInteTest, TestPkRead) { ASSERT_FALSE(query_view->expired()); ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch read_batch, reader->NextBatch()); ASSERT_FALSE(BatchReader::IsEofBatch(read_batch)); + reader->Close(); + reader.reset(); + ASSERT_TRUE(query_view->expired()); ASSERT_OK_AND_ASSIGN(std::shared_ptr read_array, ReadResultCollector::GetArray(std::move(read_batch))); ASSERT_NE(nullptr, read_array); read_array.reset(); - reader->Close(); - reader.reset(); - ASSERT_TRUE(query_view->expired()); } TEST_F(RealtimeWriteInteTest, TestPkRealtimeReadOptimizedScanUnsupported) { @@ -1142,7 +1139,7 @@ TEST_F(RealtimeWriteInteTest, TestPkDeleteInsertAndPinnedReadsAcrossRefresh) { ASSERT_OK_AND_ASSIGN(std::vector plan_rows, ReadRows(pinned_plan, realtime_context)); ASSERT_EQ((std::vector{{1, "inserted", "p0"}}), plan_rows); ASSERT_OK_AND_ASSIGN(std::shared_ptr reader_rows, - ReadResultCollector::CollectResult(pinned_reader.get())); + ReadResultCollector::CollectResult(std::move(pinned_reader))); ASSERT_EQ(1, reader_rows->length()); ASSERT_OK(writer->Close()); } @@ -1194,12 +1191,12 @@ TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, /*predicate=*/nullptr)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadPlan(plan, realtime_context, {"payload", "id"}, /*predicate=*/nullptr, /*enable_predicate_filter=*/false)); - ASSERT_NE(nullptr, result.data); - ASSERT_GT(result.data->num_chunks(), 1); - for (const std::shared_ptr& chunk : result.data->chunks()) { + ASSERT_NE(nullptr, result); + ASSERT_GT(result->num_chunks(), 1); + for (const std::shared_ptr& chunk : result->chunks()) { ASSERT_LE(chunk->length(), 2); } std::shared_ptr result_type = arrow::struct_( @@ -1214,9 +1211,8 @@ TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { [0, "disk-11", 11] ])") .ValueOrDie(); - ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) - << result.data->ToString(); - result.reader->Close(); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) + << result->ToString(); ASSERT_OK(writer->Close()); } @@ -1331,7 +1327,7 @@ TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, table_read->CreateReader(plan->Splits())); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - ReadResultCollector::CollectResult(reader.get())); + ReadResultCollector::CollectResult(std::move(reader))); const std::shared_ptr result_type = arrow::struct_({ arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("payload", arrow::struct_({projected_b})), @@ -1347,7 +1343,6 @@ TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { .ValueOrDie(); ASSERT_TRUE(std::make_shared(expected)->Equals(*actual)) << actual->ToString(); - reader->Close(); ASSERT_OK(writer->Close()); } From 99b885e176a740f5cdb9fa4c83d6f8c165aaeeb3 Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Mon, 31 Aug 2026 16:04:18 +0800 Subject: [PATCH 7/8] fix clang-tidy --- src/paimon/core/operation/merge_file_split_read.cpp | 9 ++++----- .../core/realtime/arrow_realtime_store_factory.cpp | 3 +-- src/paimon/core/table/source/append_only_table_read.cpp | 3 +-- src/paimon/core/table/source/key_value_table_read.cpp | 3 +-- 4 files changed, 7 insertions(+), 11 deletions(-) diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index 4193fb37c..55c246dc9 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -683,11 +683,10 @@ Result> MergeFileSplitRead::CreateProjectedReader( // KeyValueProjectionReader converts KeyValue objects to arrow array according to projection std::unique_ptr projection_reader; if (!context_->EnableMultiThreadRowToBatch()) { - PAIMON_ASSIGN_OR_RAISE( - projection_reader, - KeyValueProjectionReader::Create(std::move(sort_merge_reader), raw_read_schema_, - projection_, options_.GetReadBatchSize(), - arrow_pool_)); + PAIMON_ASSIGN_OR_RAISE(projection_reader, + KeyValueProjectionReader::Create( + std::move(sort_merge_reader), raw_read_schema_, projection_, + options_.GetReadBatchSize(), arrow_pool_)); } else { const int32_t thread_number = context_->GetRowToBatchThreadNumber(); assert(thread_number > 0); diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index 0edbf7953..2f7335193 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -44,8 +44,7 @@ Result> ArrowRealtimeStoreFactory::Create( arrow::ImportSchema(request.write_schema.get())); switch (request.mode) { case RealtimeStoreMode::APPEND_ONLY: { - std::shared_ptr arrow_pool = - GetSharedArrowPool(request.memory_pool); + std::shared_ptr arrow_pool = GetSharedArrowPool(request.memory_pool); return std::make_shared(imported_schema, request.statistics_mode, request.memory_pool, arrow_pool); } diff --git a/src/paimon/core/table/source/append_only_table_read.cpp b/src/paimon/core/table/source/append_only_table_read.cpp index b51ba99bd..e58946340 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -111,8 +111,7 @@ Result> AppendOnlyTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } - return std::make_unique(std::move(readers), - context_->GetArrowMemoryPool()); + return std::make_unique(std::move(readers), context_->GetArrowMemoryPool()); } Result> AppendOnlyTableRead::CreateRealtimeReader( diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 21282dd64..c962d7511 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -258,8 +258,7 @@ Result> KeyValueTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } - return std::make_unique(std::move(readers), - context_->GetArrowMemoryPool()); + return std::make_unique(std::move(readers), context_->GetArrowMemoryPool()); } Result> KeyValueTableRead::CreateRealtimeReader( From 49be4be02715389b37b6ae41d0ccb43cf6af3e94 Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Mon, 31 Aug 2026 19:00:49 +0800 Subject: [PATCH 8/8] fix comments --- include/paimon/reader/batch_reader.h | 16 ++--- ...hared_shredding_read_plan_factory_test.cpp | 6 +- .../data/shredding/shredding_file_reader.cpp | 2 +- .../apply_bitmap_index_batch_reader_test.cpp | 2 +- .../complete_index_score_batch_reader.cpp | 2 +- ...complete_index_score_batch_reader_test.cpp | 4 +- .../reader/blob_fallback_batch_reader.cpp | 2 +- .../blob_fallback_batch_reader_test.cpp | 22 +++--- .../blob_view_resolving_batch_reader.cpp | 2 +- .../blob_view_resolving_batch_reader_test.cpp | 8 +-- .../reader/complete_row_kind_batch_reader.cpp | 2 +- .../complete_row_kind_batch_reader_test.cpp | 4 +- .../reader/concat_batch_reader_test.cpp | 6 +- .../reader/data_evolution_file_reader.cpp | 3 +- .../data_evolution_file_reader_test.cpp | 20 +++--- .../late_materializing_file_batch_reader.cpp | 2 +- ...e_materializing_file_batch_reader_test.cpp | 38 +++++----- .../reader/predicate_batch_reader_test.cpp | 8 +-- .../prefetch_file_batch_reader_impl.cpp | 2 +- .../prefetch_file_batch_reader_impl_test.cpp | 71 +++++++++---------- src/paimon/common/reader/reader_utils.cpp | 3 +- .../common/reader/reader_utils_test.cpp | 6 +- src/paimon/common/utils/arrow/mem_utils.cpp | 27 +++---- src/paimon/common/utils/arrow/mem_utils.h | 11 ++- .../common/utils/arrow/mem_utils_test.cpp | 33 ++++++--- ...pply_deletion_vector_batch_reader_test.cpp | 2 +- .../global_index/global_index_write_task.cpp | 2 +- .../complete_row_tracking_fields_reader.cpp | 2 +- ...mplete_row_tracking_fields_reader_test.cpp | 10 +-- src/paimon/core/io/field_mapping_reader.cpp | 3 +- .../core/io/field_mapping_reader_test.cpp | 35 +++++---- .../io/key_value_in_memory_record_reader.h | 2 +- .../io/key_value_meta_projection_consumer.cpp | 2 +- .../io/key_value_projection_reader_test.cpp | 8 +-- .../core/io/meta_to_arrow_array_converter.cpp | 2 +- .../core/io/row_to_arrow_array_converter.h | 2 +- .../core/io/vector_file_batch_reader.cpp | 2 +- .../core/io/vector_file_batch_reader_test.cpp | 12 ++-- .../index_manifest_entry_serializer.h | 2 +- .../core/manifest/manifest_entry_serializer.h | 2 +- .../manifest/manifest_file_meta_serializer.h | 2 +- src/paimon/core/mergetree/spill_writer.h | 2 +- .../core/operation/internal_read_context.cpp | 2 +- .../operation/merge_file_split_read_test.cpp | 3 +- .../operation/raw_file_split_read_test.cpp | 6 +- .../core/postpone/postpone_bucket_writer.h | 2 +- .../core/realtime/arrow_realtime_store.cpp | 4 +- .../realtime/arrow_realtime_store_factory.cpp | 2 +- .../realtime/primary_key_realtime_store.cpp | 4 +- .../realtime/realtime_append_only_writer.cpp | 2 +- .../realtime/realtime_primary_key_writer.cpp | 4 +- .../table/source/fallback_data_split_test.cpp | 2 +- .../table/system/audit_log_system_table.cpp | 11 +-- .../table/system/in_memory_system_table.cpp | 2 +- .../core/utils/manifest_meta_reader.cpp | 3 +- src/paimon/core/utils/objects_file.h | 2 +- .../format/avro/avro_file_batch_reader.cpp | 2 +- .../format/avro/avro_file_format_test.cpp | 7 ++ .../format/avro/avro_format_writer_test.cpp | 5 +- src/paimon/format/avro/avro_reader_builder.h | 13 ++-- .../format/blob/blob_file_batch_reader.cpp | 2 +- .../blob/blob_file_batch_reader_test.cpp | 44 +++++++----- .../format/blob/blob_format_writer_test.cpp | 41 +++++------ src/paimon/format/blob/blob_reader_builder.h | 11 ++- .../format/blob/blob_stats_extractor.cpp | 3 +- .../mosaic/mosaic_file_batch_reader.cpp | 2 +- .../format/mosaic/mosaic_file_format_test.cpp | 9 ++- .../format/mosaic/mosaic_reader_builder.h | 10 +-- .../format/orc/orc_file_batch_reader_test.cpp | 7 ++ src/paimon/format/orc/orc_format_defs.h | 2 +- src/paimon/format/orc/orc_reader_builder.h | 9 ++- src/paimon/format/orc/orc_reader_wrapper.cpp | 2 +- .../parquet/parquet_file_batch_reader.cpp | 2 +- .../parquet_file_batch_reader_test.cpp | 6 ++ .../format/parquet/parquet_reader_builder.h | 11 ++- .../format/parquet/parquet_vector_io_test.cpp | 8 +-- 76 files changed, 352 insertions(+), 287 deletions(-) diff --git a/include/paimon/reader/batch_reader.h b/include/paimon/reader/batch_reader.h index 87752a64e..e6a261c6c 100644 --- a/include/paimon/reader/batch_reader.h +++ b/include/paimon/reader/batch_reader.h @@ -43,12 +43,12 @@ class PAIMON_EXPORT BatchReader { /// If EOF is reached, returns an OK status with a nullptr array. Returns an error status only /// for critical failures (e.g., IO errors). Once an error is returned, this method must not be /// retried, as it will repeatedly return the same error code. - /// \note IMPORTANT: A non-EOF ArrowArray and all its nested child arrays must have offset 0 to + /// @warning A non-EOF ArrowArray and all its nested child arrays must have offset 0 to /// avoid potential issues during conversion through the Arrow C Data Interface. - /// \note A returned ArrowArray must retain every allocator and plugin resource needed by its + /// @warning A returned ArrowArray must retain every allocator and plugin resource needed by its /// release callback, so it remains releasable after this reader is destroyed. - /// \note Consumers must treat the returned ArrowArray and ArrowSchema as one complete Arrow C - /// Data Interface ownership unit. Moving or retaining an individual child ArrowArray without + /// @warning Consumers must treat the returned ArrowArray and ArrowSchema as one complete Arrow + /// C Data Interface ownership unit. Moving or retaining an individual child ArrowArray without /// its root array is unsupported because resource lifetimes are retained by the root array's /// release chain. /// @@ -63,12 +63,12 @@ class PAIMON_EXPORT BatchReader { /// If EOF is reached, returns an OK status with a nullptr array. Returns an error status only /// for critical failures (e.g., IO errors). Once an error is returned, this method must not be /// retried, as it will repeatedly return the same error code. - /// \note IMPORTANT: A non-EOF ArrowArray and all its nested child arrays must have offset 0 to + /// @warning A non-EOF ArrowArray and all its nested child arrays must have offset 0 to /// avoid potential issues during conversion through the Arrow C Data Interface. - /// \note A returned ArrowArray must retain every allocator and plugin resource needed by its + /// @warning A returned ArrowArray must retain every allocator and plugin resource needed by its /// release callback, so it remains releasable after this reader is destroyed. - /// \note Consumers must treat the returned ArrowArray and ArrowSchema as one complete Arrow C - /// Data Interface ownership unit. Moving or retaining an individual child ArrowArray without + /// @warning Consumers must treat the returned ArrowArray and ArrowSchema as one complete Arrow + /// C Data Interface ownership unit. Moving or retaining an individual child ArrowArray without /// its root array is unsupported because resource lifetimes are retained by the root array's /// release chain. /// diff --git a/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp index a46986a12..de8ce58a6 100644 --- a/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp +++ b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp @@ -138,7 +138,7 @@ class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test { field_read_plans.emplace(field->name(), std::move(field_read_plan)); } return std::make_unique(std::move(reader), std::move(field_read_plans), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); } Result> CreateReader( @@ -303,7 +303,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjection) std::map> contexts; contexts.emplace("tags", std::move(field_read_plan)); auto reader = std::make_unique(std::move(mock_reader), std::move(contexts), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); auto read_schema = ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field})); @@ -443,7 +443,7 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjectionFr std::map> contexts; contexts.emplace("tags", std::move(field_read_plan)); auto reader = std::make_unique(std::move(mock_reader), std::move(contexts), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); auto read_schema = ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field})); diff --git a/src/paimon/common/data/shredding/shredding_file_reader.cpp b/src/paimon/common/data/shredding/shredding_file_reader.cpp index eb22ab33a..d55c2298b 100644 --- a/src/paimon/common/data/shredding/shredding_file_reader.cpp +++ b/src/paimon/common/data/shredding/shredding_file_reader.cpp @@ -109,7 +109,7 @@ Result ShreddingFileReader::NextBatchWithBitma auto new_c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*new_struct_array, new_c_array.get(), new_c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(new_c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(new_c_array.get(), new_c_schema.get(), arrow_pool_)); batch = std::make_pair(std::move(new_c_array), std::move(new_c_schema)); return batch_with_bitmap; } diff --git a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp index 1b71c137a..ef01bbeac 100644 --- a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp +++ b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp @@ -100,7 +100,7 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader.cpp b/src/paimon/common/global_index/complete_index_score_batch_reader.cpp index 65e0c72ee..75eb36104 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader.cpp +++ b/src/paimon/common/global_index/complete_index_score_batch_reader.cpp @@ -109,7 +109,7 @@ Result CompleteIndexScoreBatchReader::NextBatc arrow::StructArray::Make(array_vec, field_names_with_score_)); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*array_with_score, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return batch_with_bitmap; } } // namespace paimon diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp b/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp index 6f1495c51..ec56e0834 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp +++ b/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp @@ -52,7 +52,7 @@ class CompleteIndexScoreBatchReaderTest : public ::testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), selected_bitmap, batch_size); return std::make_unique(std::move(file_batch_reader), scores, - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); } std::unique_ptr PrepareCompleteIndexScoreBatchReader( @@ -61,7 +61,7 @@ class CompleteIndexScoreBatchReaderTest : public ::testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), batch_size); return std::make_unique(std::move(file_batch_reader), scores, - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); } private: diff --git a/src/paimon/common/reader/blob_fallback_batch_reader.cpp b/src/paimon/common/reader/blob_fallback_batch_reader.cpp index a91b04852..f3bcdd611 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader.cpp +++ b/src/paimon/common/reader/blob_fallback_batch_reader.cpp @@ -368,7 +368,7 @@ Result BlobFallbackBatchReader::NextBatch() { std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*normalized_array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return std::make_pair(std::move(c_array), std::move(c_schema)); } diff --git a/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp b/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp index 9b3fd5fe8..17ce53ded 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp +++ b/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp @@ -118,9 +118,9 @@ class BlobFallbackBatchReaderTest : public ::testing::Test { for (const auto& specs : group_specs) { groups.push_back(MakeGroup(specs, file_batch_size)); } - ASSERT_OK_AND_ASSIGN(auto reader, BlobFallbackBatchReader::Create( - std::move(groups), read_schema_, batch_size, - GetSharedArrowPool(pool_))); + ASSERT_OK_AND_ASSIGN( + auto reader, BlobFallbackBatchReader::Create(std::move(groups), read_schema_, + batch_size, GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result, paimon::test::ReadResultCollector::CollectResult( std::move(reader))); auto expected_chunk_array = std::make_shared(expected_array); @@ -292,7 +292,7 @@ TEST_F(BlobFallbackBatchReaderTest, TestRowTrackingFieldsPreserved) { ASSERT_OK_AND_ASSIGN(auto reader, BlobFallbackBatchReader::Create( std::move(groups), schema, batch_size, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN( auto result, paimon::test::ReadResultCollector::CollectResult(std::move(reader))); @@ -320,7 +320,7 @@ TEST_F(BlobFallbackBatchReaderTest, TestMisalignedGroupsFail) { groups.push_back(MakeGroup({SegmentSpec::File({"b0", "b1"})}, 1024)); ASSERT_OK_AND_ASSIGN(auto reader, BlobFallbackBatchReader::Create(std::move(groups), read_schema_, 1024, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->NextBatch(), "same number of rows"); } @@ -329,7 +329,7 @@ TEST_F(BlobFallbackBatchReaderTest, TestCreateValidation) { std::vector> single_group; single_group.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(single_group), read_schema_, 1024, - GetSharedArrowPool(pool_)), + GetArrowPool(pool_)), "at least two sequence groups"); // the read schema must contain a blob field @@ -338,16 +338,16 @@ TEST_F(BlobFallbackBatchReaderTest, TestCreateValidation) { groups.push_back(MakeGroup({SegmentSpec::File({"b1"})}, 1024)); auto plain_schema = arrow::schema({arrow::field("not_blob", arrow::large_binary(), /*nullable=*/true)}); - ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(groups), plain_schema, 1024, - GetSharedArrowPool(pool_)), - "should contain a blob field"); + ASSERT_NOK_WITH_MSG( + BlobFallbackBatchReader::Create(std::move(groups), plain_schema, 1024, GetArrowPool(pool_)), + "should contain a blob field"); // groups must not be empty std::vector> with_empty_group; with_empty_group.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); with_empty_group.emplace_back(); ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(with_empty_group), read_schema_, - 1024, GetSharedArrowPool(pool_)), + 1024, GetArrowPool(pool_)), "should not be empty"); // a gap segment must cover at least one selected row id @@ -355,7 +355,7 @@ TEST_F(BlobFallbackBatchReaderTest, TestCreateValidation) { with_empty_gap.push_back(MakeGroup({SegmentSpec::File({"b0"})}, 1024)); with_empty_gap.push_back(MakeGroup({SegmentSpec::GapRanges({})}, 1024)); ASSERT_NOK_WITH_MSG(BlobFallbackBatchReader::Create(std::move(with_empty_gap), read_schema_, - 1024, GetSharedArrowPool(pool_)), + 1024, GetArrowPool(pool_)), "at least one selected row id"); } diff --git a/src/paimon/common/reader/blob_view_resolving_batch_reader.cpp b/src/paimon/common/reader/blob_view_resolving_batch_reader.cpp index 822e202d1..37d05020b 100644 --- a/src/paimon/common/reader/blob_view_resolving_batch_reader.cpp +++ b/src/paimon/common/reader/blob_view_resolving_batch_reader.cpp @@ -90,7 +90,7 @@ Result BlobViewResolvingBatchReader::NextBatch() { arrow::StructArray::Make(new_fields, field_names)); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*resolved_struct_array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return batch; } diff --git a/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp b/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp index bd3f9c291..bcb60c0c7 100644 --- a/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp +++ b/src/paimon/common/reader/blob_view_resolving_batch_reader_test.cpp @@ -139,7 +139,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestEofBatch) { return std::shared_ptr(); }); BlobViewResolvingBatchReader reader(std::move(inner_reader), {"blob_col"}, std::move(resolver), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto batch, reader.NextBatch()); ASSERT_TRUE(BatchReader::IsEofBatch(batch)); } @@ -158,7 +158,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestEmptyReadBlobViewFields) { auto inner_reader = std::make_unique(struct_array); auto reader = std::make_unique( std::move(inner_reader), /*read_blob_view_fields=*/std::vector(), - std::move(resolver), GetSharedArrowPool(pool_)); + std::move(resolver), GetArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto expected_array = std::make_shared(struct_array); ASSERT_TRUE(expected_array->Equals(*result_array)); @@ -190,7 +190,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestResolvesBlobViewColumn) { auto inner_reader = std::make_unique(src_struct); auto reader = std::make_unique( std::move(inner_reader), std::vector{"blob_col"}, std::move(resolver), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto struct_array = std::dynamic_pointer_cast(result_array->chunk(0)); @@ -210,7 +210,7 @@ TEST_F(BlobViewResolvingBatchReaderTest, TestResolverError) { }); auto inner_reader = std::make_unique(src_struct); BlobViewResolvingBatchReader reader(std::move(inner_reader), {"blob_col"}, std::move(resolver), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); ASSERT_NOK_WITH_MSG(reader.NextBatch(), "cache miss"); } diff --git a/src/paimon/common/reader/complete_row_kind_batch_reader.cpp b/src/paimon/common/reader/complete_row_kind_batch_reader.cpp index f7d8e2ac2..c9591df8f 100644 --- a/src/paimon/common/reader/complete_row_kind_batch_reader.cpp +++ b/src/paimon/common/reader/complete_row_kind_batch_reader.cpp @@ -81,7 +81,7 @@ Result CompleteRowKindBatchReader::NextBatchWi arrow::StructArray::Make(fields_with_row_kind, field_names_with_row_kind_)); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*array_with_row_kind, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return batch_with_bitmap; } diff --git a/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp b/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp index 7ae10be5c..a2229ba54 100644 --- a/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp +++ b/src/paimon/common/reader/complete_row_kind_batch_reader_test.cpp @@ -69,7 +69,7 @@ class CompleteRowKindBatchReaderTest : public ::testing::Test { EXPECT_OK(orc_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt)); return std::make_unique(std::move(orc_batch_reader), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); } std::unique_ptr PrepareCompleteRowKindBatchReader( @@ -77,7 +77,7 @@ class CompleteRowKindBatchReaderTest : public ::testing::Test { auto file_batch_reader = std::make_unique(src_array, src_array->type(), batch_size); return std::make_unique(std::move(file_batch_reader), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); } private: diff --git a/src/paimon/common/reader/concat_batch_reader_test.cpp b/src/paimon/common/reader/concat_batch_reader_test.cpp index 654066743..3cc66e085 100644 --- a/src/paimon/common/reader/concat_batch_reader_test.cpp +++ b/src/paimon/common/reader/concat_batch_reader_test.cpp @@ -133,7 +133,7 @@ class ConcatBatchReaderTest : public ::testing::Test { readers.push_back(std::move(reader)); } auto concat_reader = - std::make_unique(std::move(readers), GetSharedArrowPool(pool_)); + std::make_unique(std::move(readers), GetArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_chunk_array, ReadResultCollector::CollectResult(std::move(concat_reader))); if (expected.empty()) { @@ -249,7 +249,7 @@ TEST_F(ConcatBatchReaderTest, TestReleaseReaderAtEof) { std::make_unique(data, data->type(), /*read_batch_size=*/1)); lifetime.reset(); - ConcatBatchReader reader(std::move(readers), GetSharedArrowPool(pool_)); + ConcatBatchReader reader(std::move(readers), GetArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch, reader.NextBatchWithBitmap()); ASSERT_FALSE(BatchReader::IsEofBatch(batch)); ASSERT_TRUE(weak_lifetime.expired()); @@ -271,7 +271,7 @@ TEST_F(ConcatBatchReaderTest, TestCollectMetricsAfterReleasingReaders) { std::make_unique(data2, type, /*read_batch_size=*/1), /*latency=*/17, /*io_count=*/3)); auto concat_reader = - std::make_unique(std::move(readers), GetSharedArrowPool(pool_)); + std::make_unique(std::move(readers), GetArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(std::shared_ptr result, ReadResultCollector::CollectResult(concat_reader.get())); diff --git a/src/paimon/common/reader/data_evolution_file_reader.cpp b/src/paimon/common/reader/data_evolution_file_reader.cpp index 55e3aabb9..347fe4392 100644 --- a/src/paimon/common/reader/data_evolution_file_reader.cpp +++ b/src/paimon/common/reader/data_evolution_file_reader.cpp @@ -103,7 +103,8 @@ Result DataEvolutionFileReader::NextBatchWithB std::unique_ptr target_c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*target_array, target_c_arrow_array.get(), target_c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(target_c_arrow_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK( + AddArrowArrayLifetime(target_c_arrow_array.get(), target_c_schema.get(), arrow_pool_)); auto target_batch = std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema)); return ReaderUtils::AddAllValidBitmap(std::move(target_batch)); } diff --git a/src/paimon/common/reader/data_evolution_file_reader_test.cpp b/src/paimon/common/reader/data_evolution_file_reader_test.cpp index 924ceadd3..35cdda40b 100644 --- a/src/paimon/common/reader/data_evolution_file_reader_test.cpp +++ b/src/paimon/common/reader/data_evolution_file_reader_test.cpp @@ -77,7 +77,7 @@ class DataEvolutionFileReaderTest : public ::testing::Test, ASSERT_OK_AND_ASSIGN(auto data_evolution_file_reader, DataEvolutionFileReader::Create( std::move(readers), read_schema, batch_size, reader_offsets, - field_offsets, GetSharedArrowPool(pool_))); + field_offsets, GetArrowPool(pool_))); // check metrics, data_evolution_file_reader collects all row of each // MockFileBatchReader auto metrics = data_evolution_file_reader->GetReaderMetrics(); @@ -111,7 +111,7 @@ class DataEvolutionFileReaderTest : public ::testing::Test, readers.push_back(std::move(file_batch_reader)); DataEvolutionFileReader fake_data_evolution_reader( std::move(readers), /*read_schema=*/arrow::schema({}), read_batch_size, - /*reader_offsets=*/{}, /*field_offsets=*/{}, GetSharedArrowPool(pool_)); + /*reader_offsets=*/{}, /*field_offsets=*/{}, GetArrowPool(pool_)); arrow::ArrayVector result_array_vec; while (true) { ASSERT_OK_AND_ASSIGN(auto result_array, @@ -148,7 +148,7 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { arrow::FieldVector read_fields; auto read_schema = arrow::schema(read_fields); ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, - {}, {}, GetSharedArrowPool(pool_)), + {}, {}, GetArrowPool(pool_)), "read schema must not be empty"); } { @@ -163,7 +163,7 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { std::vector field_offsets = {0, 1, 0}; ASSERT_NOK_WITH_MSG( DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, reader_offsets, - field_offsets, GetSharedArrowPool(pool_)), + field_offsets, GetArrowPool(pool_)), "read schema, row offsets and field offsets must have the same size"); } { @@ -178,7 +178,7 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { std::vector field_offsets = {0, 1, 1, 0}; ASSERT_NOK_WITH_MSG( DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, reader_offsets, - field_offsets, GetSharedArrowPool(pool_)), + field_offsets, GetArrowPool(pool_)), "readers must not be empty"); } { @@ -194,10 +194,10 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) { auto read_schema = arrow::schema(read_fields); std::vector reader_offsets = {0, 0, 1, 1}; std::vector field_offsets = {0, 1, 1, 0}; - ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create( - std::move(readers), read_schema, /*read_batch_size=*/10, - reader_offsets, field_offsets, GetSharedArrowPool(pool_)), - "reader offset is out of range of readers"); + ASSERT_NOK_WITH_MSG( + DataEvolutionFileReader::Create(std::move(readers), read_schema, /*read_batch_size=*/10, + reader_offsets, field_offsets, GetArrowPool(pool_)), + "reader offset is out of range of readers"); } } @@ -638,7 +638,7 @@ TEST_P(DataEvolutionFileReaderTest, TestSingleReaderRowCountMismatch) { ASSERT_OK_AND_ASSIGN( auto data_evolution_file_reader, DataEvolutionFileReader::Create(std::move(readers), read_schema, /*read_batch_size=*/10, - reader_offsets, field_offsets, GetSharedArrowPool(pool_))); + reader_offsets, field_offsets, GetArrowPool(pool_))); // array0 has 6 rows but array1 only has 5 rows ASSERT_NOK_WITH_MSG( paimon::test::ReadResultCollector::CollectResult(std::move(data_evolution_file_reader)), diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 6fab25e04..23b6960fd 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -237,7 +237,7 @@ Result LateMaterializingFileBatchReader::AssembleFul std::unique_ptr<::ArrowSchema> c_schema = std::make_unique<::ArrowSchema>(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*full_struct, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return std::make_pair(std::move(c_array), std::move(c_schema)); } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index f3aac3e93..c86635f50 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -227,7 +227,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr, std::nullopt)); @@ -246,7 +246,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); // read schema is just {k}; the predicate on k covers all columns -> payload empty auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(0l)); @@ -267,7 +267,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ContiguousSubsetAcrossBatches) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(4l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -292,7 +292,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ScatteredAlternatingMatch) { auto data = BuildData(ks); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -316,7 +316,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { auto data = BuildData(ks); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); // predicate hits {1,3,5,7,9,11}; selection keeps only {1,5,9} @@ -342,7 +342,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(100l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -356,7 +356,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, SeekToRowRealignsProbeCursor) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(5l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -385,7 +385,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); auto* mock_ptr = mock.get(); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(2l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -414,7 +414,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); auto predicate1 = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(7l)); @@ -445,7 +445,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ForwardsRowCountAndFileSchema) { auto data = BuildData({0, 1, 2, 3}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(uint64_t num_rows, reader->GetNumberOfRows()); EXPECT_EQ(num_rows, 4u); @@ -462,7 +462,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { auto type = data->type(); auto mock = std::make_unique(data, type, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); // probe columns = {a (idx0), c (idx2)}; payload columns = {b, d, e} auto pred_a = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, Literal(3l)); @@ -506,7 +506,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { auto type = data->type(); auto mock = std::make_unique(data, type, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); // probe = {k}; payload = {arr (list), tag} auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); @@ -539,7 +539,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); LateMaterializingReaderBuilder builder( std::make_unique(data, full_type_, /*batch_size=*/3), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); ASSERT_OK_AND_ASSIGN( @@ -549,7 +549,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l)); ::ArrowSchema c_schema; @@ -575,7 +575,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); LateMaterializingReaderBuilder builder( std::make_unique(data, full_type_, /*batch_size=*/3), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); ASSERT_OK_AND_ASSIGN( @@ -585,7 +585,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto full_schema = arrow::schema(full_fields_); auto predicate1 = @@ -638,7 +638,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee // and each range-honoring reader only reads its assigned slice. LateMaterializingReaderBuilder builder( std::make_unique(data, full_type_, /*batch_size=*/3), - GetSharedArrowPool(pool_)); + GetArrowPool(pool_)); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(3)); ASSERT_OK_AND_ASSIGN( @@ -648,7 +648,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee /*prefetch_max_parallel_num=*/3, /*batch_size=*/3, /*prefetch_batch_count=*/6, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); ::ArrowSchema c_schema; @@ -675,7 +675,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, FailsOnPredicateTypeMismatch) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create( - std::move(mock), GetSharedArrowPool(pool_))); + std::move(mock), GetArrowPool(pool_))); // k is int64 in the schema, but the predicate claims FieldType::INT (int32). auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::INT, Literal(10)); diff --git a/src/paimon/common/reader/predicate_batch_reader_test.cpp b/src/paimon/common/reader/predicate_batch_reader_test.cpp index 3b6c70c9a..a47159a8c 100644 --- a/src/paimon/common/reader/predicate_batch_reader_test.cpp +++ b/src/paimon/common/reader/predicate_batch_reader_test.cpp @@ -81,7 +81,7 @@ class PredicateBatchReaderTest : public ::testing::Test { const std::shared_ptr& expected_array) const { ASSERT_OK_AND_ASSIGN(auto predicate_reader, PredicateBatchReader::Create(std::move(reader), predicate, - GetSharedArrowPool(GetDefaultPool()))); + GetArrowPool(GetDefaultPool()))); ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, ReadResultCollector::CollectResult(std::move(predicate_reader))); if (expected_array) { @@ -205,9 +205,9 @@ TEST_F(PredicateBatchReaderTest, TestFullAndEmptyCase) { TEST_F(PredicateBatchReaderTest, TestInvalidInput) { auto data_array = PrepareArray(8); auto reader = std::make_unique(data_array, data_type_, /*batch_size=*/10); - ASSERT_NOK_WITH_MSG(PredicateBatchReader::Create(std::move(reader), nullptr, - GetSharedArrowPool(GetDefaultPool())), - "create predicate batch reader failed. predicate is nullptr"); + ASSERT_NOK_WITH_MSG( + PredicateBatchReader::Create(std::move(reader), nullptr, GetArrowPool(GetDefaultPool())), + "create predicate batch reader failed. predicate is nullptr"); } } // namespace paimon::test diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp index aaa2cd6e4..f24dd2b5a 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -654,7 +654,7 @@ Status PrefetchFileBatchReaderImpl::HandleReadResult( ArrowUtils::NormalizeArrayOffsets(sliced_array, arrow_pool_.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*normalized_array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); bitmap.RemoveRange(slice_end, array->length()); global_row_ids = std::vector(global_row_ids.begin(), global_row_ids.begin() + slice_end); diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp index d408b8641..02f9ba8d4 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp @@ -269,7 +269,7 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, read_ahead_cache_enabled, CacheConfig(), - /*enable_io_metrics=*/true, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/true, pool_, GetArrowPool(pool_))); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); EXPECT_TRUE(arrow_status.ok()); @@ -360,7 +360,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestSimple) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); if (prefetch_max_parallel_num == 1) { ASSERT_NOK( reader->GetReaderMetrics()->GetCounter(PrefetchIoMetrics::READ_LATENCY_COUNT)); @@ -388,7 +388,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLimits) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/true, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/true, pool_, GetArrowPool(pool_))); // simulate read limits, only read 8 batches for (int32_t i = 0; i < 8; i++) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, @@ -442,7 +442,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithoutInitializeReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); // simulate read limits, only read 8 batches ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "prefetch reader read ranges are not initialized"); @@ -460,7 +460,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestFailedIoMetrics) { /*prefetch_max_parallel_num=*/1, /*batch_size=*/10, /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/false, CacheConfig(), - /*enable_io_metrics=*/true, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/true, pool_, GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "injected synchronous read failure"); std::shared_ptr metrics = reader->GetReaderMetrics(); @@ -549,7 +549,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_OK(prefetch_reader->RefreshReadRanges()); std::vector> read_ranges_0 = {{0, 30}, {90, 101}}; @@ -580,7 +580,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRangesDisablePrefetchByAdapti /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/true, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); ASSERT_FALSE(reader->NeedPrefetch()); std::shared_ptr metrics = reader->GetReaderMetrics(); @@ -605,7 +605,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRanges) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_FALSE(prefetch_reader->need_prefetch_); prefetch_reader->need_prefetch_ = true; @@ -649,7 +649,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRangesReturnErrorWhenPushDownFail prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->need_prefetch_ = true; @@ -676,7 +676,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, WorkloopSetReadStatusWhenCacheInitFailed prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, - invalid_cache_config, /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + invalid_cache_config, /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->Workloop(); @@ -697,7 +697,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenShutdown) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->is_shutdown_ = true; @@ -716,7 +716,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenNoCurrentReadRang prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->read_ranges_in_group_ = {{}}; @@ -735,7 +735,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLargeBatchSize) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -756,7 +756,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPartialReaderSuccessRead) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { dynamic_cast(prefetch_reader->readers_[i].get()) @@ -802,7 +802,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { @@ -841,7 +841,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithEmptyData) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -861,7 +861,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCallNextBatchAfterReadingEof) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -905,7 +905,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCreateReaderWithoutNextBatch) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); } TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { @@ -920,7 +920,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*prefetch_max_parallel_num=*/0, batch_size, 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -928,7 +928,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, /*batch_size=*/-1, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -937,7 +937,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*enable_adaptive_prefetch_strategy=*/false, /*executor=*/nullptr, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), /*enable_io_metrics=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -945,7 +945,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -953,7 +953,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /*fs=*/nullptr, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); } { ASSERT_OK_AND_ASSIGN( @@ -963,7 +963,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SeekToRow(/*row_number=*/101), "not support seek to row for prefetch reader"); } @@ -1060,22 +1060,21 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithBitmap) { MockFormatReaderBuilder reader_builder(data_array, data_type_, bitmap, /*read_batch_size=*/100); int32_t prefetch_max_parallel_num = 3; - ASSERT_OK_AND_ASSIGN(auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, - /*batch_size=*/100, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, - /*read_ahead_cache_enabled=*/true, CacheConfig(), - /*enable_io_metrics=*/false, pool_, GetSharedArrowPool(pool_))); + ASSERT_OK_AND_ASSIGN(auto reader, PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, + &reader_builder, mock_fs_, prefetch_max_parallel_num, + /*batch_size=*/100, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, + /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, pool_, GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_chunk_array, ReadResultCollector::CollectResult(std::move(reader))); ASSERT_OK_AND_ASSIGN(auto data_batch, ReadResultCollector::GetReadBatch(data_array)); - ASSERT_OK_AND_ASSIGN(auto expected_batch, ReaderUtils::ApplyBitmapToReadBatch( - std::make_pair(std::move(data_batch), bitmap), - GetSharedArrowPool(pool_))); + ASSERT_OK_AND_ASSIGN(auto expected_batch, + ReaderUtils::ApplyBitmapToReadBatch( + std::make_pair(std::move(data_batch), bitmap), GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto expected_array, ReadResultCollector::GetArray(std::move(expected_batch))); auto expected_chunk_array = std::make_shared(expected_array); diff --git a/src/paimon/common/reader/reader_utils.cpp b/src/paimon/common/reader/reader_utils.cpp index 1d7b166de..0fe38d2c3 100644 --- a/src/paimon/common/reader/reader_utils.cpp +++ b/src/paimon/common/reader/reader_utils.cpp @@ -100,7 +100,8 @@ Result ReaderUtils::ApplyBitmapToReadBatch( std::unique_ptr result_c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*result, result_c_array.get(), result_c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(result_c_array.get(), arrow_pool)); + PAIMON_RETURN_NOT_OK( + AddArrowArrayLifetime(result_c_array.get(), result_c_schema.get(), arrow_pool)); return make_pair(std::move(result_c_array), std::move(result_c_schema)); } diff --git a/src/paimon/common/reader/reader_utils_test.cpp b/src/paimon/common/reader/reader_utils_test.cpp index ad0ee3b05..b9848987a 100644 --- a/src/paimon/common/reader/reader_utils_test.cpp +++ b/src/paimon/common/reader/reader_utils_test.cpp @@ -90,9 +90,9 @@ TEST(ReaderUtilsTest, TestCollectResultReleasesBufferedBatchesOnError) { auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[1]").ValueOrDie(); ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, ReadResultCollector::GetReadBatch(array)); - std::shared_ptr lifetime = std::make_shared(1); - std::weak_ptr weak_lifetime = lifetime; - ASSERT_OK(AddArrowArrayLifetime(batch.first.get(), lifetime)); + std::shared_ptr lifetime = std::make_shared(1); + std::weak_ptr weak_lifetime = lifetime; + ASSERT_OK(AddArrowArrayLifetime(batch.first.get(), batch.second.get(), lifetime)); lifetime.reset(); auto reader = std::make_unique(std::move(batch)); diff --git a/src/paimon/common/utils/arrow/mem_utils.cpp b/src/paimon/common/utils/arrow/mem_utils.cpp index 99eab4ce5..348cf4d01 100644 --- a/src/paimon/common/utils/arrow/mem_utils.cpp +++ b/src/paimon/common/utils/arrow/mem_utils.cpp @@ -126,20 +126,20 @@ class ArrowMemPoolAdaptor : public arrow::MemoryPool { } // namespace -std::unique_ptr GetArrowPool(const std::shared_ptr& pool) { - return std::make_unique(pool); -} - -std::shared_ptr GetSharedArrowPool(const std::shared_ptr& pool) { +std::shared_ptr GetArrowPool(const std::shared_ptr& pool) { return std::make_shared(pool); } -Status AddArrowArrayLifetime(ArrowArray* array, const std::shared_ptr& lifetime) { +Status AddArrowArrayLifetime(ArrowArray* array, ArrowSchema* schema, + const std::shared_ptr& lifetime) { if (array == nullptr || array->release == nullptr) { return Status::Invalid("cannot add lifetime to a released ArrowArray"); } if (lifetime.use_count() == 0) { ArrowArrayRelease(array); + if (schema != nullptr && schema->release != nullptr) { + ArrowSchemaRelease(schema); + } return Status::Invalid("cannot add an empty lifetime to an ArrowArray"); } if (array->release == ReleaseArrowArray) { @@ -148,16 +148,11 @@ Status AddArrowArrayLifetime(ArrowArray* array, const std::shared_ptr& lif return Status::OK(); } } - try { - std::unique_ptr data = std::make_unique( - ArrowArrayPrivateData{array->release, array->private_data, lifetime}); - array->release = ReleaseArrowArray; - array->private_data = data.release(); - return Status::OK(); - } catch (const std::bad_alloc&) { - ArrowArrayRelease(array); - return Status::OutOfMemory("failed to add lifetime to an ArrowArray"); - } + std::unique_ptr data = std::make_unique( + ArrowArrayPrivateData{array->release, array->private_data, lifetime}); + array->release = ReleaseArrowArray; + array->private_data = data.release(); + return Status::OK(); } } // namespace paimon diff --git a/src/paimon/common/utils/arrow/mem_utils.h b/src/paimon/common/utils/arrow/mem_utils.h index 101245b57..e229606b7 100644 --- a/src/paimon/common/utils/arrow/mem_utils.h +++ b/src/paimon/common/utils/arrow/mem_utils.h @@ -28,13 +28,11 @@ #include "paimon/visibility.h" struct ArrowArray; +struct ArrowSchema; namespace paimon { -PAIMON_EXPORT std::unique_ptr GetArrowPool( - const std::shared_ptr& pool); - -PAIMON_EXPORT std::shared_ptr GetSharedArrowPool( +PAIMON_EXPORT std::shared_ptr GetArrowPool( const std::shared_ptr& pool); /// Keep an additional resource alive until the ArrowArray is released. @@ -42,8 +40,9 @@ PAIMON_EXPORT std::shared_ptr GetSharedArrowPool( /// The existing release callback and private data are preserved as a release chain, so this helper /// can be applied to arrays produced by Arrow libraries or format plugins without inspecting their /// private data. Once a valid, unreleased array is passed in, this function releases it if -/// retaining the lifetime fails. -PAIMON_EXPORT Status AddArrowArrayLifetime(ArrowArray* array, +/// retaining the lifetime fails. If a paired schema was exported, pass it so the schema is also +/// released on failure; otherwise, pass nullptr. +PAIMON_EXPORT Status AddArrowArrayLifetime(ArrowArray* array, ArrowSchema* schema, const std::shared_ptr& lifetime); } // namespace paimon diff --git a/src/paimon/common/utils/arrow/mem_utils_test.cpp b/src/paimon/common/utils/arrow/mem_utils_test.cpp index 530572169..30c7afd44 100644 --- a/src/paimon/common/utils/arrow/mem_utils_test.cpp +++ b/src/paimon/common/utils/arrow/mem_utils_test.cpp @@ -222,8 +222,8 @@ TEST(MemUtilsTest, TestAddArrowArrayLifetimeComposesReleaseChain) { array.release = ReleaseTrackingArrowArray; array.private_data = new TrackingArrowArrayPrivateData{original_lifetime, &release_order}; - ASSERT_OK(AddArrowArrayLifetime(&array, inner_lifetime)); - ASSERT_OK(AddArrowArrayLifetime(&array, outer_lifetime)); + ASSERT_OK(AddArrowArrayLifetime(&array, /*schema=*/nullptr, inner_lifetime)); + ASSERT_OK(AddArrowArrayLifetime(&array, /*schema=*/nullptr, outer_lifetime)); original_lifetime.reset(); inner_lifetime.reset(); outer_lifetime.reset(); @@ -239,7 +239,7 @@ TEST(MemUtilsTest, TestAddArrowArrayLifetimeKeepsPoolAliveUntilOriginalRelease) std::shared_ptr paimon_pool = std::make_shared(&pool_destroyed, &free_count); std::weak_ptr weak_paimon_pool = paimon_pool; - std::shared_ptr arrow_pool = GetSharedArrowPool(paimon_pool); + std::shared_ptr arrow_pool = GetArrowPool(paimon_pool); ArrowArray c_array{}; { arrow::Int32Builder builder(arrow_pool.get()); @@ -249,7 +249,7 @@ TEST(MemUtilsTest, TestAddArrowArrayLifetimeKeepsPoolAliveUntilOriginalRelease) std::shared_ptr array = std::move(array_result).ValueOrDie(); ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); } - ASSERT_OK(AddArrowArrayLifetime(&c_array, arrow_pool)); + ASSERT_OK(AddArrowArrayLifetime(&c_array, /*schema=*/nullptr, arrow_pool)); arrow_pool.reset(); paimon_pool.reset(); @@ -270,9 +270,9 @@ TEST(MemUtilsTest, TestAddArrowArrayLifetimeDeduplicatesSameOwner) { array.release = ReleaseTrackingArrowArray; array.private_data = new TrackingArrowArrayPrivateData{nullptr, &release_order}; - ASSERT_OK(AddArrowArrayLifetime(&array, lifetime)); + ASSERT_OK(AddArrowArrayLifetime(&array, /*schema=*/nullptr, lifetime)); const int64_t use_count = lifetime.use_count(); - ASSERT_OK(AddArrowArrayLifetime(&array, lifetime)); + ASSERT_OK(AddArrowArrayLifetime(&array, /*schema=*/nullptr, lifetime)); ASSERT_EQ(use_count, lifetime.use_count()); ArrowArrayRelease(&array); @@ -280,15 +280,30 @@ TEST(MemUtilsTest, TestAddArrowArrayLifetimeDeduplicatesSameOwner) { TEST(MemUtilsTest, TestAddArrowArrayLifetimeRejectsInvalidInput) { ArrowArray array{}; - ASSERT_NOK(AddArrowArrayLifetime(nullptr, std::make_shared(1))); - ASSERT_NOK(AddArrowArrayLifetime(&array, std::make_shared(1))); + ASSERT_NOK(AddArrowArrayLifetime(nullptr, /*schema=*/nullptr, std::make_shared(1))); + ASSERT_NOK(AddArrowArrayLifetime(&array, /*schema=*/nullptr, std::make_shared(1))); std::vector release_order; array.release = ReleaseTrackingArrowArray; array.private_data = new TrackingArrowArrayPrivateData{nullptr, &release_order}; - ASSERT_NOK(AddArrowArrayLifetime(&array, nullptr)); + ASSERT_NOK(AddArrowArrayLifetime(&array, /*schema=*/nullptr, nullptr)); ASSERT_EQ(nullptr, array.release); ASSERT_EQ(std::vector({0}), release_order); } +TEST(MemUtilsTest, TestAddArrowArrayLifetimeReleasesPairedSchemaOnFailure) { + std::vector release_order; + ArrowArray array{}; + array.release = ReleaseTrackingArrowArray; + array.private_data = new TrackingArrowArrayPrivateData{nullptr, &release_order}; + ArrowSchema schema{}; + ASSERT_TRUE( + arrow::ExportSchema(*arrow::schema({arrow::field("value", arrow::int32())}), &schema).ok()); + + ASSERT_NOK(AddArrowArrayLifetime(&array, &schema, nullptr)); + ASSERT_EQ(nullptr, array.release); + ASSERT_EQ(nullptr, schema.release); + ASSERT_EQ(std::vector({0}), release_order); +} + } // namespace paimon::test diff --git a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp index a1baa2868..5f4ef074f 100644 --- a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp +++ b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp @@ -70,7 +70,7 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, void CheckResult(const std::string& data_str, const std::vector& dv_data, const std::string& expected_str) { std::shared_ptr pool = GetDefaultPool(); - std::shared_ptr arrow_pool = GetSharedArrowPool(pool); + std::shared_ptr arrow_pool = GetArrowPool(pool); auto f1 = arrow::ipc::internal::json::ArrayFromJSON(int_type_, data_str).ValueOrDie(); std::shared_ptr data = arrow::StructArray::Make({f1}, target_type_->fields()).ValueOrDie(); diff --git a/src/paimon/core/global_index/global_index_write_task.cpp b/src/paimon/core/global_index/global_index_write_task.cpp index 6ff929554..0288d1b96 100644 --- a/src/paimon/core/global_index/global_index_write_task.cpp +++ b/src/paimon/core/global_index/global_index_write_task.cpp @@ -357,7 +357,7 @@ Result> GlobalIndexWriteTask::WriteIndex( } const auto& range = ranges[0]; std::shared_ptr pool = memory_pool ? memory_pool : GetDefaultPool(); - std::unique_ptr arrow_pool = GetArrowPool(pool); + std::shared_ptr arrow_pool = GetArrowPool(pool); // load schema PAIMON_ASSIGN_OR_RAISE(CoreOptions tmp_options, CoreOptions::FromMap(options, file_system)); diff --git a/src/paimon/core/io/complete_row_tracking_fields_reader.cpp b/src/paimon/core/io/complete_row_tracking_fields_reader.cpp index 3bd347677..e9d703bc8 100644 --- a/src/paimon/core/io/complete_row_tracking_fields_reader.cpp +++ b/src/paimon/core/io/complete_row_tracking_fields_reader.cpp @@ -149,7 +149,7 @@ CompleteRowTrackingFieldsBatchReader::NextBatchWithBitmap() { arrow::StructArray::Make(sub_array_vec, read_schema_->field_names())); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*target_array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return std::move(src_array_with_bitmap); } diff --git a/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp b/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp index d070d2cbb..806080b15 100644 --- a/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp +++ b/src/paimon/core/io/complete_row_tracking_fields_reader_test.cpp @@ -51,7 +51,7 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { auto complete_row_tracking_fields_reader = std::make_unique( std::move(file_batch_reader), first_row_id, snapshot_id, - /*file_field_names=*/std::nullopt, GetSharedArrowPool(pool_)); + /*file_field_names=*/std::nullopt, GetArrowPool(pool_)); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK(complete_row_tracking_fields_reader->SetReadSchema( @@ -76,7 +76,7 @@ class CompleteRowTrackingFieldsBatchReaderTest : public testing::Test { auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/10, /*snapshot_id=*/1, - file_field_names, GetSharedArrowPool(pool_)); + file_field_names, GetArrowPool(pool_)); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( @@ -378,7 +378,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidWithReadNonExistFiel auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/100, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, GetSharedArrowPool(pool_)); + /*file_field_names=*/std::nullopt, GetArrowPool(pool_)); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( @@ -415,7 +415,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidNextBatchBeforeSetRe auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/100, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, GetSharedArrowPool(pool_)); + /*file_field_names=*/std::nullopt, GetArrowPool(pool_)); ASSERT_NOK_WITH_MSG(complete_row_tracking_fields_reader->NextBatchWithBitmap(), "in CompleteRowTrackingFieldsBatchReader SetReadSchema is supposed to be " "called before NextBatch"); @@ -446,7 +446,7 @@ TEST_F(CompleteRowTrackingFieldsBatchReaderTest, TestInvalidNullFirstRowId) { auto complete_row_tracking_fields_reader = std::make_shared( std::move(file_batch_reader), /*first_row_id=*/std::nullopt, /*snapshot_id=*/8, - /*file_field_names=*/std::nullopt, GetSharedArrowPool(pool_)); + /*file_field_names=*/std::nullopt, GetArrowPool(pool_)); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*read_schema, &c_read_schema).ok()); ASSERT_OK( diff --git a/src/paimon/core/io/field_mapping_reader.cpp b/src/paimon/core/io/field_mapping_reader.cpp index 3a0986393..68a3135e3 100644 --- a/src/paimon/core/io/field_mapping_reader.cpp +++ b/src/paimon/core/io/field_mapping_reader.cpp @@ -321,7 +321,8 @@ Result FieldMappingReader::NextBatchWithBitmap std::unique_ptr target_c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*arrow_array, target_c_arrow_array.get(), target_c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(target_c_arrow_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK( + AddArrowArrayLifetime(target_c_arrow_array.get(), target_c_schema.get(), arrow_pool_)); auto target_batch = std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema)); return std::make_pair(std::move(target_batch), std::move(bitmap)); } diff --git a/src/paimon/core/io/field_mapping_reader_test.cpp b/src/paimon/core/io/field_mapping_reader_test.cpp index 19ac1c537..139b7fb20 100644 --- a/src/paimon/core/io/field_mapping_reader_test.cpp +++ b/src/paimon/core/io/field_mapping_reader_test.cpp @@ -137,11 +137,10 @@ class FieldMappingReaderTest : public ::testing::Test { /*batch_size=*/1); ASSERT_OK_AND_ASSIGN( - auto reader, - FieldMappingReader::Create( - /*field_count=*/read_schema->num_fields(), std::move(orc_batch_reader), partition_, - std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, - GetSharedArrowPool(pool_))); + auto reader, FieldMappingReader::Create( + /*field_count=*/read_schema->num_fields(), std::move(orc_batch_reader), + partition_, std::move(mapping), + /*skip_map_selected_keys_filter_field_ids=*/{}, GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); if (expect_array == nullptr && result_array == nullptr) { @@ -180,11 +179,10 @@ class FieldMappingReaderTest : public ::testing::Test { /*predicate=*/mapping->non_partition_info.non_partition_filter, /*batch_size=*/1); ASSERT_OK_AND_ASSIGN( - auto reader, - FieldMappingReader::Create( - /*field_count=*/read_schema->num_fields(), std::move(orc_batch_reader), partition, - std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, - GetSharedArrowPool(pool_))); + auto reader, FieldMappingReader::Create( + /*field_count=*/read_schema->num_fields(), std::move(orc_batch_reader), + partition, std::move(mapping), + /*skip_map_selected_keys_filter_field_ids=*/{}, GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); if (expect_array == nullptr && result_array == nullptr) { @@ -247,7 +245,7 @@ TEST_F(FieldMappingReaderTest, TestGenerateSinglePartitionArray) { auto mapping_reader, FieldMappingReader::Create( /*field_count=*/8, /*reader=*/nullptr, partition, std::move(field_mapping), - /*skip_map_selected_keys_filter_field_ids=*/{}, GetSharedArrowPool(pool_))); + /*skip_map_selected_keys_filter_field_ids=*/{}, GetArrowPool(pool_))); { ASSERT_OK_AND_ASSIGN(auto p7_array, mapping_reader->GenerateSinglePartitionArray( @@ -479,7 +477,7 @@ TEST_F(FieldMappingReaderTest, TestSchemaEvolutionAddedFieldInsideList) { FieldMappingReader::Create(read_schema->num_fields(), std::move(mock), BinaryRow::EmptyRow(), std::move(mapping), /*skip_map_selected_keys_filter_field_ids=*/{}, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(reader))); auto expect_array = @@ -865,12 +863,11 @@ TEST_F(FieldMappingReaderTest, TestCreateFailFastOnInvalidMapSelectedKeysMetadat /*partition_info=*/PartitionInfo(), std::move(non_part_info), /*non_exist_field_info=*/std::nullopt); - ASSERT_NOK_WITH_MSG( - FieldMappingReader::Create( - /*field_count=*/1, - /*reader=*/nullptr, - /*partition=*/BinaryRow::EmptyRow(), std::move(field_mapping), - /*skip_map_selected_keys_filter_field_ids=*/{}, GetSharedArrowPool(pool_)), - "Duplicate selected key 'a'"); + ASSERT_NOK_WITH_MSG(FieldMappingReader::Create( + /*field_count=*/1, + /*reader=*/nullptr, + /*partition=*/BinaryRow::EmptyRow(), std::move(field_mapping), + /*skip_map_selected_keys_filter_field_ids=*/{}, GetArrowPool(pool_)), + "Duplicate selected key 'a'"); } } // namespace paimon::test diff --git a/src/paimon/core/io/key_value_in_memory_record_reader.h b/src/paimon/core/io/key_value_in_memory_record_reader.h index 169446544..8ffd0c7ba 100644 --- a/src/paimon/core/io/key_value_in_memory_record_reader.h +++ b/src/paimon/core/io/key_value_in_memory_record_reader.h @@ -85,7 +85,7 @@ class KeyValueInMemoryRecordReader : public KeyValueRecordReader { std::vector user_defined_sequence_fields_; bool sequence_fields_ascending_ = true; std::shared_ptr pool_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::shared_ptr value_struct_array_; std::vector row_kinds_; std::shared_ptr key_comparator_; diff --git a/src/paimon/core/io/key_value_meta_projection_consumer.cpp b/src/paimon/core/io/key_value_meta_projection_consumer.cpp index 7e9117412..259aad6d0 100644 --- a/src/paimon/core/io/key_value_meta_projection_consumer.cpp +++ b/src/paimon/core/io/key_value_meta_projection_consumer.cpp @@ -60,7 +60,7 @@ Result> KeyValueMetaProjectionCo target_to_src_mapping.size())); } - std::shared_ptr arrow_pool = GetSharedArrowPool(pool); + std::shared_ptr arrow_pool = GetArrowPool(pool); // target fields of output array: special fields + value fields std::unique_ptr array_builder; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder( diff --git a/src/paimon/core/io/key_value_projection_reader_test.cpp b/src/paimon/core/io/key_value_projection_reader_test.cpp index 1b15a77e1..f25ce1f20 100644 --- a/src/paimon/core/io/key_value_projection_reader_test.cpp +++ b/src/paimon/core/io/key_value_projection_reader_test.cpp @@ -96,14 +96,14 @@ class KeyValueProjectionReaderTest : public testing::Test, if (!multi_thread_row_to_batch) { EXPECT_OK_AND_ASSIGN( auto projection_reader, - KeyValueProjectionReader::Create( - std::move(sort_merge_reader), target_schema, target_to_src_mapping, - /*batch_size=*/batch_size, GetSharedArrowPool(pool_))); + KeyValueProjectionReader::Create(std::move(sort_merge_reader), target_schema, + target_to_src_mapping, + /*batch_size=*/batch_size, GetArrowPool(pool_))); return std::move(projection_reader); } else { return std::make_unique( std::move(sort_merge_reader), target_schema, target_to_src_mapping, batch_size, - /*projection_thread_num=*/3, GetSharedArrowPool(pool_)); + /*projection_thread_num=*/3, GetArrowPool(pool_)); } } diff --git a/src/paimon/core/io/meta_to_arrow_array_converter.cpp b/src/paimon/core/io/meta_to_arrow_array_converter.cpp index 1980a3478..2aa9a75c8 100644 --- a/src/paimon/core/io/meta_to_arrow_array_converter.cpp +++ b/src/paimon/core/io/meta_to_arrow_array_converter.cpp @@ -27,7 +27,7 @@ Result> MetaToArrowArrayConverter::Cr return Status::Invalid("meta_data_type in MetaToArrowArrayConverter must be struct type"); } auto struct_type = checked_pointer_cast(meta_data_type); - std::shared_ptr arrow_pool = GetSharedArrowPool(pool); + std::shared_ptr arrow_pool = GetArrowPool(pool); std::unique_ptr array_builder; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder( arrow_pool.get(), arrow::struct_(struct_type->fields()), &array_builder)); diff --git a/src/paimon/core/io/row_to_arrow_array_converter.h b/src/paimon/core/io/row_to_arrow_array_converter.h index b4153d968..6caa9f618 100644 --- a/src/paimon/core/io/row_to_arrow_array_converter.h +++ b/src/paimon/core/io/row_to_arrow_array_converter.h @@ -110,7 +110,7 @@ Result RowToArrowArrayConverter::FinishAndAccumula std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return make_pair(std::move(c_array), std::move(c_schema)); } diff --git a/src/paimon/core/io/vector_file_batch_reader.cpp b/src/paimon/core/io/vector_file_batch_reader.cpp index 518958307..da475c9ae 100644 --- a/src/paimon/core/io/vector_file_batch_reader.cpp +++ b/src/paimon/core/io/vector_file_batch_reader.cpp @@ -264,7 +264,7 @@ Result VectorFileBatchReader::ConvertBatch(ReadBatch&& b PAIMON_ASSIGN_OR_RAISE(array, ConvertToReadType(array, read_type_, arrow_pool_.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW(array->Validate()); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return std::move(batch); } diff --git a/src/paimon/core/io/vector_file_batch_reader_test.cpp b/src/paimon/core/io/vector_file_batch_reader_test.cpp index ec86a1487..560ea5b22 100644 --- a/src/paimon/core/io/vector_file_batch_reader_test.cpp +++ b/src/paimon/core/io/vector_file_batch_reader_test.cpp @@ -62,7 +62,7 @@ TEST(VectorFileBatchReaderTest, ConvertSchemaAndNextBatch) { std::make_unique(physical_array, physical_type, /*batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); MockFileBatchReader* inner_reader = mock_reader.get(); - VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); + VectorFileBatchReader reader(std::move(mock_reader), GetArrowPool(GetDefaultPool())); ASSERT_TRUE(VectorFileBatchReader::ContainsVector(arrow::schema(logical_type->fields()))); ASSERT_FALSE(VectorFileBatchReader::ContainsVector(arrow::schema(physical_type->fields()))); @@ -99,7 +99,7 @@ TEST(VectorFileBatchReaderTest, KeepFixedSizeListFileSchema) { mock_reader->EnableRandomizeBatchSize(false); MockFileBatchReader* inner_reader = mock_reader.get(); auto reader = std::make_unique(std::move(mock_reader), - GetSharedArrowPool(GetDefaultPool())); + GetArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); @@ -139,7 +139,7 @@ TEST(VectorFileBatchReaderTest, NormalizeFixedSizeListElementField) { std::make_unique(file_array, file_type, /*batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); auto reader = std::make_unique(std::move(mock_reader), - GetSharedArrowPool(GetDefaultPool())); + GetArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); @@ -179,7 +179,7 @@ TEST(VectorFileBatchReaderTest, ConvertNestedVectorsWithBitmap) { /*read_batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); auto reader = std::make_unique(std::move(mock_reader), - GetSharedArrowPool(GetDefaultPool())); + GetArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_read_schema, /*predicate=*/nullptr, @@ -209,7 +209,7 @@ TEST(VectorFileBatchReaderTest, RejectInvalidVectorValues) { auto mock_reader = std::make_unique(physical_array, physical_type, /*read_batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); - VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); + VectorFileBatchReader reader(std::move(mock_reader), GetArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE( arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); @@ -228,7 +228,7 @@ TEST(VectorFileBatchReaderTest, RejectInvalidFixedSizeListVectorValues) { auto mock_reader = std::make_unique(physical_array, physical_type, /*read_batch_size=*/10); mock_reader->EnableRandomizeBatchSize(false); - VectorFileBatchReader reader(std::move(mock_reader), GetSharedArrowPool(GetDefaultPool())); + VectorFileBatchReader reader(std::move(mock_reader), GetArrowPool(GetDefaultPool())); ArrowSchema c_read_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(physical_type->fields()), &c_read_schema).ok()); ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, diff --git a/src/paimon/core/manifest/index_manifest_entry_serializer.h b/src/paimon/core/manifest/index_manifest_entry_serializer.h index dd19f515e..904db20f9 100644 --- a/src/paimon/core/manifest/index_manifest_entry_serializer.h +++ b/src/paimon/core/manifest/index_manifest_entry_serializer.h @@ -50,6 +50,6 @@ class IndexManifestEntrySerializer : public VersionedObjectSerializer arrow_pool_; + std::shared_ptr arrow_pool_; }; } // namespace paimon diff --git a/src/paimon/core/manifest/manifest_entry_serializer.h b/src/paimon/core/manifest/manifest_entry_serializer.h index 4a71a1b68..9d4ddcfa7 100644 --- a/src/paimon/core/manifest/manifest_entry_serializer.h +++ b/src/paimon/core/manifest/manifest_entry_serializer.h @@ -63,7 +63,7 @@ class ManifestEntrySerializer : public VersionedObjectSerializer private: static constexpr int32_t VERSION_1 = 1; static constexpr int32_t VERSION_2 = 2; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; DataFileMetaSerializer data_file_meta_serializer_; }; } // namespace paimon diff --git a/src/paimon/core/manifest/manifest_file_meta_serializer.h b/src/paimon/core/manifest/manifest_file_meta_serializer.h index b7260656c..7067d38f6 100644 --- a/src/paimon/core/manifest/manifest_file_meta_serializer.h +++ b/src/paimon/core/manifest/manifest_file_meta_serializer.h @@ -56,7 +56,7 @@ class ManifestFileMetaSerializer : public VersionedObjectSerializer arrow_pool_; + std::shared_ptr arrow_pool_; }; } // namespace paimon diff --git a/src/paimon/core/mergetree/spill_writer.h b/src/paimon/core/mergetree/spill_writer.h index 6676dc284..3aaf1b046 100644 --- a/src/paimon/core/mergetree/spill_writer.h +++ b/src/paimon/core/mergetree/spill_writer.h @@ -73,7 +73,7 @@ class SpillWriter { bool use_threads_; std::shared_ptr out_stream_; std::shared_ptr arrow_output_stream_adapter_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::shared_ptr arrow_writer_; FileIOChannel::ID channel_id_; bool closed_ = false; diff --git a/src/paimon/core/operation/internal_read_context.cpp b/src/paimon/core/operation/internal_read_context.cpp index 3858929e7..dd60ce404 100644 --- a/src/paimon/core/operation/internal_read_context.cpp +++ b/src/paimon/core/operation/internal_read_context.cpp @@ -285,7 +285,7 @@ Result> InternalReadContext::Create( if (!context->GetMemoryPool()) { return Status::Invalid("memory pool is null"); } - std::shared_ptr arrow_pool = GetSharedArrowPool(context->GetMemoryPool()); + std::shared_ptr arrow_pool = GetArrowPool(context->GetMemoryPool()); return std::unique_ptr( new InternalReadContext(context, table_schema, read_schema, core_options, arrow_pool)); } diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index 02e4387d9..2096c2aa9 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -363,8 +363,7 @@ class MergeFileSplitReadTest : public ::testing::Test, split_read->CreateReader(split)); batch_readers.emplace_back(std::move(reader)); } - return std::make_unique(std::move(batch_readers), - GetSharedArrowPool(pool_)); + return std::make_unique(std::move(batch_readers), GetArrowPool(pool_)); } private: diff --git a/src/paimon/core/operation/raw_file_split_read_test.cpp b/src/paimon/core/operation/raw_file_split_read_test.cpp index 82b132991..90528004b 100644 --- a/src/paimon/core/operation/raw_file_split_read_test.cpp +++ b/src/paimon/core/operation/raw_file_split_read_test.cpp @@ -173,8 +173,8 @@ class RawFileSplitReadTest : public ::testing::Test { split_read->CreateReader(split)); batch_readers.emplace_back(std::move(reader)); } - auto batch_reader = std::make_unique(std::move(batch_readers), - GetSharedArrowPool(pool_)); + auto batch_reader = + std::make_unique(std::move(batch_readers), GetArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_TRUE(result_array->Equals(expected_array)); @@ -425,7 +425,7 @@ TEST_F(RawFileSplitReadTest, TestEmptyPlan) { ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, split_read->CreateReader(data_split)); batch_readers.push_back(std::move(reader)); auto batch_reader = - std::make_unique(std::move(batch_readers), GetSharedArrowPool(pool_)); + std::make_unique(std::move(batch_readers), GetArrowPool(pool_)); ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(std::move(batch_reader))); ASSERT_EQ(result_array, nullptr); diff --git a/src/paimon/core/postpone/postpone_bucket_writer.h b/src/paimon/core/postpone/postpone_bucket_writer.h index f61c98876..53e00af27 100644 --- a/src/paimon/core/postpone/postpone_bucket_writer.h +++ b/src/paimon/core/postpone/postpone_bucket_writer.h @@ -133,7 +133,7 @@ class PostponeBucketWriter : public BatchWriter { private: std::shared_ptr pool_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; std::vector trimmed_primary_keys_; CoreOptions options_; std::shared_ptr path_factory_; diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp index 6dc9bbca2..bad0d0706 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store.cpp @@ -155,7 +155,7 @@ class ArrowRealtimeStore::CommitBatchReader : public BatchReader { auto c_array = std::make_unique(); auto c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*result, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return ReadBatch(std::move(c_array), std::move(c_schema)); } @@ -221,7 +221,7 @@ class ArrowRealtimeStore::QueryBatchReader : public BatchReader { auto c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*output, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return ReadBatchWithBitmap(ReadBatch(std::move(c_array), std::move(c_schema)), std::move(candidate_rows)); } diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index 2f7335193..dff12b589 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -44,7 +44,7 @@ Result> ArrowRealtimeStoreFactory::Create( arrow::ImportSchema(request.write_schema.get())); switch (request.mode) { case RealtimeStoreMode::APPEND_ONLY: { - std::shared_ptr arrow_pool = GetSharedArrowPool(request.memory_pool); + std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); return std::make_shared(imported_schema, request.statistics_mode, request.memory_pool, arrow_pool); } diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index 3e46e0788..2e8a4185c 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -112,7 +112,7 @@ class StoredBatchReader final : public BatchReader { ArrowUtils::NormalizeRecordBatchOffsets(record_batch, arrow_pool_.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportRecordBatch(*normalized_batch, array.get(), schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(array.get(), schema.get(), arrow_pool_)); data_.reset(); arrow_pool_.reset(); export_guard.Release(); @@ -276,7 +276,7 @@ Result> PrimaryKeyRealtimeStore::Create if (!memory_pool) { return Status::Invalid("PK real-time store memory pool is null"); } - std::shared_ptr arrow_pool = GetSharedArrowPool(memory_pool); + std::shared_ptr arrow_pool = GetArrowPool(memory_pool); return std::shared_ptr(new PrimaryKeyRealtimeStore( std::make_unique(transport_schema, std::move(arrow_pool)))); } diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 0733d9e93..78d570a54 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -70,7 +70,7 @@ RealtimeAppendOnlyWriter::RealtimeAppendOnlyWriter( const std::shared_ptr& file_writer, const std::shared_ptr& input_schema, int64_t next_offset, const std::shared_ptr& memory_pool) - : arrow_pool_(GetSharedArrowPool(memory_pool)), + : arrow_pool_(GetArrowPool(memory_pool)), realtime_store_(realtime_store), file_writer_(file_writer), input_schema_(input_schema), diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index b2c04c4c5..0b2849a63 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -158,7 +158,7 @@ RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( int64_t next_offset, int64_t last_sequence_number, const std::shared_ptr& memory_pool) : memory_pool_(memory_pool), - arrow_pool_(GetSharedArrowPool(memory_pool)), + arrow_pool_(GetArrowPool(memory_pool)), realtime_store_(realtime_store), merge_tree_writer_(merge_tree_writer), realtime_context_(realtime_context), @@ -205,7 +205,7 @@ Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { arrow_pool_.get())); auto output = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*transport, output.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(output.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(output.get(), /*schema=*/nullptr, arrow_pool_)); RecordBatchBuilder builder(output.get()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr transport_batch, builder.Finish()); PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ diff --git a/src/paimon/core/table/source/fallback_data_split_test.cpp b/src/paimon/core/table/source/fallback_data_split_test.cpp index e37a7e190..11607f392 100644 --- a/src/paimon/core/table/source/fallback_data_split_test.cpp +++ b/src/paimon/core/table/source/fallback_data_split_test.cpp @@ -72,7 +72,7 @@ TEST(FallbackTableReadTest, RoutesIndexedSplitToMainTable) { TrackingTableRead* main_table_ptr = main_table.get(); TrackingTableRead* fallback_table_ptr = fallback_table.get(); FallbackTableRead table_read(std::move(main_table), std::move(fallback_table), - GetSharedArrowPool(pool)); + GetArrowPool(pool)); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/"", /*data_files=*/{}); 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 7e35a39ca..3b5f1c824 100644 --- a/src/paimon/core/table/system/audit_log_system_table.cpp +++ b/src/paimon/core/table/system/audit_log_system_table.cpp @@ -137,7 +137,8 @@ class ChangelogBatchReader : public BatchReader { 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())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(output_c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK( + AddArrowArrayLifetime(output_c_array.get(), output_c_schema.get(), arrow_pool_)); return std::make_pair(std::move(output_c_array), std::move(output_c_schema)); } } @@ -293,7 +294,7 @@ class ChangelogTableRead : public TableRead { output_schema_(std::move(output_schema)), include_sequence_number_(include_sequence_number), converter_(std::move(converter)), - arrow_pool_(GetSharedArrowPool(pool)) {} + arrow_pool_(GetArrowPool(pool)) {} Result> CreateReader( const std::vector>& splits) override { @@ -349,9 +350,9 @@ 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, GetSharedArrowPool(pool)); + return std::make_unique(std::move(reader), std::move(output_schema), + include_sequence_number, std::move(converter), + pack_update_before_after, GetArrowPool(pool)); } AuditLogSystemTable::AuditLogSystemTable(std::shared_ptr fs, std::string table_path, diff --git a/src/paimon/core/table/system/in_memory_system_table.cpp b/src/paimon/core/table/system/in_memory_system_table.cpp index c05a012bb..8d6111d79 100644 --- a/src/paimon/core/table/system/in_memory_system_table.cpp +++ b/src/paimon/core/table/system/in_memory_system_table.cpp @@ -76,7 +76,7 @@ class InMemorySystemTableRead : public TableRead { public: InMemorySystemTableRead(std::shared_ptr table, const std::shared_ptr& memory_pool) - : table_(std::move(table)), arrow_pool_(GetSharedArrowPool(memory_pool)) {} + : table_(std::move(table)), arrow_pool_(GetArrowPool(memory_pool)) {} Result> CreateReader( const std::vector>& splits) override { diff --git a/src/paimon/core/utils/manifest_meta_reader.cpp b/src/paimon/core/utils/manifest_meta_reader.cpp index 3fca1156d..acb67cb8a 100644 --- a/src/paimon/core/utils/manifest_meta_reader.cpp +++ b/src/paimon/core/utils/manifest_meta_reader.cpp @@ -62,7 +62,8 @@ Result ManifestMetaReader::NextBatch() { PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*target_array, target_c_arrow_array.get(), target_c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(target_c_arrow_array.get(), pool_)); + PAIMON_RETURN_NOT_OK( + AddArrowArrayLifetime(target_c_arrow_array.get(), target_c_schema.get(), pool_)); return std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema)); } diff --git a/src/paimon/core/utils/objects_file.h b/src/paimon/core/utils/objects_file.h index 09274d5e3..b79450433 100644 --- a/src/paimon/core/utils/objects_file.h +++ b/src/paimon/core/utils/objects_file.h @@ -109,7 +109,7 @@ ObjectsFile::ObjectsFile(const std::shared_ptr& file_system, const std::shared_ptr& pool) : path_factory_(path_factory), pool_(pool), - arrow_pool_(GetSharedArrowPool(pool)), + arrow_pool_(GetArrowPool(pool)), serializer_(std::move(serializer)), writer_builder_(std::move(writer_builder)), file_system_(file_system), diff --git a/src/paimon/format/avro/avro_file_batch_reader.cpp b/src/paimon/format/avro/avro_file_batch_reader.cpp index 7f1048dd5..252f563c3 100644 --- a/src/paimon/format/avro/avro_file_batch_reader.cpp +++ b/src/paimon/format/avro/avro_file_batch_reader.cpp @@ -133,7 +133,7 @@ Result AvroFileBatchReader::NextBatch() { std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); previous_batch_row_count_ = c_array->length; return make_pair(std::move(c_array), std::move(c_schema)); } catch (const ::avro::Exception& e) { diff --git a/src/paimon/format/avro/avro_file_format_test.cpp b/src/paimon/format/avro/avro_file_format_test.cpp index fe361940a..71727ad44 100644 --- a/src/paimon/format/avro/avro_file_format_test.cpp +++ b/src/paimon/format/avro/avro_file_format_test.cpp @@ -29,6 +29,7 @@ #include "gtest/gtest.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/path_util.h" +#include "paimon/format/avro/avro_reader_builder.h" #include "paimon/format/file_format.h" #include "paimon/format/file_format_factory.h" #include "paimon/format/format_writer.h" @@ -102,6 +103,12 @@ class AvroFileFormatTest : public testing::Test, public ::testing::WithParamInte std::unique_ptr dir_; }; +TEST(AvroReaderBuilderTest, RejectsNullMemoryPool) { + AvroReaderBuilder builder(/*batch_size=*/1024); + builder.WithMemoryPool(nullptr); + ASSERT_NOK_WITH_MSG(builder.Build(nullptr), "Avro reader memory pool is nullptr"); +} + INSTANTIATE_TEST_SUITE_P(Compression, AvroFileFormatTest, ::testing::ValuesIn(std::vector( {"zstd", "zstandard", "snappy", "null", "deflate"}))); diff --git a/src/paimon/format/avro/avro_format_writer_test.cpp b/src/paimon/format/avro/avro_format_writer_test.cpp index 099d532c8..571e99c07 100644 --- a/src/paimon/format/avro/avro_format_writer_test.cpp +++ b/src/paimon/format/avro/avro_format_writer_test.cpp @@ -119,9 +119,8 @@ class AvroFormatWriterTest : public ::testing::Test { void CheckResult(const std::string& file_path, int32_t row_count) const { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs_->Open(file_path)); - ASSERT_OK_AND_ASSIGN( - auto file_reader, - AvroFileBatchReader::Create(input_stream, 1024, pool_, GetSharedArrowPool(pool_))); + ASSERT_OK_AND_ASSIGN(auto file_reader, AvroFileBatchReader::Create( + input_stream, 1024, pool_, GetArrowPool(pool_))); ASSERT_OK_AND_ASSIGN(uint64_t num_rows, file_reader->GetNumberOfRows()); ASSERT_EQ(num_rows, row_count); diff --git a/src/paimon/format/avro/avro_reader_builder.h b/src/paimon/format/avro/avro_reader_builder.h index 4d450d4ae..977a347a2 100644 --- a/src/paimon/format/avro/avro_reader_builder.h +++ b/src/paimon/format/avro/avro_reader_builder.h @@ -33,18 +33,23 @@ namespace paimon::avro { class AvroReaderBuilder : public ReaderBuilder { public: explicit AvroReaderBuilder(int32_t batch_size) - : batch_size_(batch_size), - pool_(GetDefaultPool()), - arrow_pool_(GetSharedArrowPool(pool_)) {} + : batch_size_(batch_size), pool_(GetDefaultPool()), arrow_pool_(GetArrowPool(pool_)) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; - arrow_pool_ = GetSharedArrowPool(pool); + if (pool == nullptr) { + arrow_pool_.reset(); + } else { + arrow_pool_ = GetArrowPool(pool); + } return this; } Result> Build( const std::shared_ptr& path) const override { + if (pool_ == nullptr) { + return Status::Invalid("Avro reader memory pool is nullptr"); + } return AvroFileBatchReader::Create(path, batch_size_, pool_, arrow_pool_); } diff --git a/src/paimon/format/blob/blob_file_batch_reader.cpp b/src/paimon/format/blob/blob_file_batch_reader.cpp index 8f7b0649d..f5e3ef9ce 100644 --- a/src/paimon/format/blob/blob_file_batch_reader.cpp +++ b/src/paimon/format/blob/blob_file_batch_reader.cpp @@ -325,7 +325,7 @@ Result BlobFileBatchReader::NextBatch() { std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*blob_array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); previous_batch_start_pos_ = current_pos_; current_pos_ += rows_to_read; previous_batch_row_count_ = c_array->length; diff --git a/src/paimon/format/blob/blob_file_batch_reader_test.cpp b/src/paimon/format/blob/blob_file_batch_reader_test.cpp index 3b66b6ee8..17482af77 100644 --- a/src/paimon/format/blob/blob_file_batch_reader_test.cpp +++ b/src/paimon/format/blob/blob_file_batch_reader_test.cpp @@ -25,6 +25,7 @@ #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/data/blob.h" #include "paimon/format/blob/blob_format_writer.h" +#include "paimon/format/blob/blob_reader_builder.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/read_result_collector.h" @@ -32,6 +33,12 @@ namespace paimon::blob::test { +TEST(BlobReaderBuilderTest, RejectsNullMemoryPool) { + BlobReaderBuilder builder(/*batch_size=*/10, /*options=*/{}); + builder.WithMemoryPool(nullptr); + ASSERT_NOK_WITH_MSG(builder.Build(nullptr), "Blob reader memory pool is nullptr"); +} + class BlobFileBatchReaderTest : public testing::Test, public ::testing::WithParamInterface { public: void SetUp() override { @@ -47,10 +54,10 @@ class BlobFileBatchReaderTest : public testing::Test, public ::testing::WithPara std::shared_ptr fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs->Open(table_path + "/bucket-0/" + paimon_blob_file)); - ASSERT_OK_AND_ASSIGN(auto reader, BlobFileBatchReader::Create( - input_stream, /*batch_size=*/1024, blob_as_descriptor, - /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + ASSERT_OK_AND_ASSIGN(auto reader, + BlobFileBatchReader::Create( + input_stream, /*batch_size=*/1024, blob_as_descriptor, + /*emit_placeholder_sentinel=*/false, pool_, GetArrowPool(pool_))); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, selection_bitmap)); ASSERT_OK_AND_ASSIGN(auto chunked_array, paimon::test::ReadResultCollector::CollectResult(std::move(reader))); @@ -170,7 +177,7 @@ TEST_F(BlobFileBatchReaderTest, TestRowNumbers) { BlobFileBatchReader::Create(input_stream, /*batch_size=*/1, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, std::nullopt)); ASSERT_OK_AND_ASSIGN(auto number_of_rows, reader->GetNumberOfRows()); @@ -213,7 +220,7 @@ TEST_F(BlobFileBatchReaderTest, TestRowNumbersWithSelectionBitmap) { BlobFileBatchReader::Create(input_stream, /*batch_size=*/1, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); RoaringBitmap32 selection; selection.Add(0); @@ -250,23 +257,22 @@ TEST_F(BlobFileBatchReaderTest, InvalidScenario) { BlobFileBatchReader::Create(input_stream, /*batch_size=*/0, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_)), + GetArrowPool(pool_)), "blob file batch reader create failed: read batch size '0' should be larger than zero"); } { - ASSERT_NOK_WITH_MSG( - BlobFileBatchReader::Create(/*input_stream=*/nullptr, - /*batch_size=*/1, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_)), - "blob file batch reader create failed: input stream is nullptr"); + ASSERT_NOK_WITH_MSG(BlobFileBatchReader::Create( + /*input_stream=*/nullptr, + /*batch_size=*/1, /*blob_as_descriptor=*/true, + /*emit_placeholder_sentinel=*/false, pool_, GetArrowPool(pool_)), + "blob file batch reader create failed: input stream is nullptr"); } { ASSERT_OK_AND_ASSIGN( auto reader, BlobFileBatchReader::Create(/*input_stream=*/input_stream, /*batch_size=*/1, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->GetFileSchema(), "blob file has no self-describing file schema"); ASSERT_TRUE(reader->GetReaderMetrics()); @@ -304,7 +310,7 @@ TEST_P(BlobFileBatchReaderTest, EmptyFile) { BlobFileBatchReader::Create(input_stream, /*batch_size=*/1, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_OK(reader->SetReadSchema(&c_schema, nullptr, std::nullopt)); ASSERT_OK_AND_ASSIGN(auto number_of_rows, reader->GetNumberOfRows()); @@ -330,7 +336,7 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { auto reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(/*read_schema=*/nullptr, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "SetReadSchema failed: read schema cannot be nullptr"); @@ -356,7 +362,7 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { auto reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "read schema field number 2 is not 1"); @@ -382,7 +388,7 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { auto reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ASSERT_NOK_WITH_MSG(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt), "field my_blob_field: large_binary is not BLOB"); @@ -406,7 +412,7 @@ TEST_F(BlobFileBatchReaderTest, SetReadSchemaWithInvalidInputs) { auto reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); RoaringBitmap32 roaring; roaring.Add(0); roaring.Add(1); diff --git a/src/paimon/format/blob/blob_format_writer_test.cpp b/src/paimon/format/blob/blob_format_writer_test.cpp index 5341900db..b0edbe3be 100644 --- a/src/paimon/format/blob/blob_format_writer_test.cpp +++ b/src/paimon/format/blob/blob_format_writer_test.cpp @@ -177,7 +177,7 @@ class BlobFormatWriterTestBase : public ::testing::Test { BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, /*emit_placeholder_sentinel=*/false, - pool_, GetSharedArrowPool(pool_))); + pool_, GetArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); @@ -265,11 +265,10 @@ TEST_P(BlobFormatWriterTest, TestSimple) { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); ASSERT_TRUE(input_stream); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + BlobFileBatchReader::Create( + input_stream, /*batch_size=*/1024, blob_as_descriptor_, + /*emit_placeholder_sentinel=*/false, pool_, GetArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -463,11 +462,10 @@ TEST_P(BlobFormatWriterTest, TestLargeBlob) { // Verify we can read it back ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + BlobFileBatchReader::Create( + input_stream, /*batch_size=*/1024, blob_as_descriptor_, + /*emit_placeholder_sentinel=*/false, pool_, GetArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -514,11 +512,10 @@ TEST_P(BlobFormatWriterTest, TestAddBatchWithNullValues) { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(dir_->Str() + "/file.blob")); ASSERT_TRUE(input_stream); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, blob_as_descriptor_, - /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + BlobFileBatchReader::Create( + input_stream, /*batch_size=*/1024, blob_as_descriptor_, + /*emit_placeholder_sentinel=*/false, pool_, GetArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -1046,7 +1043,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, @@ -1062,7 +1059,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, /*emit_placeholder_sentinel=*/true, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, @@ -1086,7 +1083,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderStrictAndAwareModes) BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/true, /*emit_placeholder_sentinel=*/true, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, @@ -1117,7 +1114,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestReadPlaceholderWithSelectionBitmap) BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, /*emit_placeholder_sentinel=*/true, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -1150,7 +1147,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestSentinelBytesVerbatimWithoutPlacehol BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, /*emit_placeholder_sentinel=*/false, pool_, - GetSharedArrowPool(pool_))); + GetArrowPool(pool_))); auto schema = arrow::schema(struct_type_->fields()); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); @@ -1187,7 +1184,7 @@ TEST_F(BlobFormatWriterPlaceholderTest, TestSentinelPrefixedValueVerbatimInPlace std::unique_ptr reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/false, emit_placeholder_sentinel, - pool_, GetSharedArrowPool(pool_))); + pool_, GetArrowPool(pool_))); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK(reader->SetReadSchema(&c_schema, /*predicate=*/nullptr, diff --git a/src/paimon/format/blob/blob_reader_builder.h b/src/paimon/format/blob/blob_reader_builder.h index c66465922..9f4460dc3 100644 --- a/src/paimon/format/blob/blob_reader_builder.h +++ b/src/paimon/format/blob/blob_reader_builder.h @@ -36,17 +36,24 @@ class BlobReaderBuilder : public ReaderBuilder { BlobReaderBuilder(int32_t batch_size, const std::map& options) : batch_size_(batch_size), pool_(GetDefaultPool()), - arrow_pool_(GetSharedArrowPool(pool_)), + arrow_pool_(GetArrowPool(pool_)), options_(options) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; - arrow_pool_ = GetSharedArrowPool(pool); + if (pool == nullptr) { + arrow_pool_.reset(); + } else { + arrow_pool_ = GetArrowPool(pool); + } return this; } Result> Build( const std::shared_ptr& input_stream) const override { + if (pool_ == nullptr) { + return Status::Invalid("Blob reader memory pool is nullptr"); + } PAIMON_ASSIGN_OR_RAISE( bool blob_as_descriptor, OptionsUtils::GetValueFromMap(options_, Options::BLOB_AS_DESCRIPTOR, false)); diff --git a/src/paimon/format/blob/blob_stats_extractor.cpp b/src/paimon/format/blob/blob_stats_extractor.cpp index 8cbbafe09..f4bc757fa 100644 --- a/src/paimon/format/blob/blob_stats_extractor.cpp +++ b/src/paimon/format/blob/blob_stats_extractor.cpp @@ -57,8 +57,7 @@ BlobStatsExtractor::ExtractWithFileInfo(const std::shared_ptr& file_ std::unique_ptr blob_reader, BlobFileBatchReader::Create(input_stream, /*batch_size=*/1024, /*blob_as_descriptor=*/true, - /*emit_placeholder_sentinel=*/false, pool, - GetSharedArrowPool(pool))); + /*emit_placeholder_sentinel=*/false, pool, GetArrowPool(pool))); ColumnStatsVector result_stats; result_stats.push_back( ColumnStats::CreateStringColumnStats(std::nullopt, std::nullopt, std::nullopt)); diff --git a/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp b/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp index 097b04584..ed508a954 100644 --- a/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp +++ b/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp @@ -199,7 +199,7 @@ Result MosaicFileBatchReader::NextBatch() { auto ffi_schema = std::make_unique<::ArrowSchema>(); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*normalized_array, ffi_array.get(), ffi_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(ffi_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(ffi_array.get(), ffi_schema.get(), arrow_pool_)); return std::make_pair(std::move(ffi_array), std::move(ffi_schema)); } diff --git a/src/paimon/format/mosaic/mosaic_file_format_test.cpp b/src/paimon/format/mosaic/mosaic_file_format_test.cpp index deb951191..3646b06ff 100644 --- a/src/paimon/format/mosaic/mosaic_file_format_test.cpp +++ b/src/paimon/format/mosaic/mosaic_file_format_test.cpp @@ -182,7 +182,7 @@ class MosaicFileFormatTest : public ::testing::Test { std::shared_ptr file_system_; std::unique_ptr directory_; std::shared_ptr pool_; - std::unique_ptr arrow_pool_; + std::shared_ptr arrow_pool_; }; TEST_F(MosaicFileFormatTest, WriteThenRead) { @@ -200,6 +200,13 @@ TEST_F(MosaicFileFormatTest, WriteThenRead) { AssertReadWithBatchSizes(path, schema, expected, {1, 2, 3, 5, 8}); } +TEST_F(MosaicFileFormatTest, ReaderBuilderRejectsNullMemoryPool) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader_builder, + format_->CreateReaderBuilder(/*batch_size=*/2)); + reader_builder->WithMemoryPool(nullptr); + ASSERT_NOK_WITH_MSG(reader_builder->Build(nullptr), "Mosaic reader memory pool is nullptr"); +} + TEST_F(MosaicFileFormatTest, EmptyProjectionPreservesRowCount) { std::string path = PathUtil::JoinPath(directory_->Str(), "empty-projection.mosaic"); arrow::FieldVector fields = {arrow::field("id", arrow::int32(), false), diff --git a/src/paimon/format/mosaic/mosaic_reader_builder.h b/src/paimon/format/mosaic/mosaic_reader_builder.h index 34bc11eb9..94da74cbf 100644 --- a/src/paimon/format/mosaic/mosaic_reader_builder.h +++ b/src/paimon/format/mosaic/mosaic_reader_builder.h @@ -31,13 +31,15 @@ namespace paimon::mosaic { class MosaicReaderBuilder : public ReaderBuilder { public: explicit MosaicReaderBuilder(int32_t batch_size) - : batch_size_(batch_size), - pool_(GetDefaultPool()), - arrow_pool_(GetSharedArrowPool(pool_)) {} + : batch_size_(batch_size), pool_(GetDefaultPool()), arrow_pool_(GetArrowPool(pool_)) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; - arrow_pool_ = GetSharedArrowPool(pool); + if (pool == nullptr) { + arrow_pool_.reset(); + } else { + arrow_pool_ = GetArrowPool(pool); + } return this; } diff --git a/src/paimon/format/orc/orc_file_batch_reader_test.cpp b/src/paimon/format/orc/orc_file_batch_reader_test.cpp index 048cdc0ff..3f57f1fbf 100644 --- a/src/paimon/format/orc/orc_file_batch_reader_test.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader_test.cpp @@ -40,6 +40,7 @@ #include "paimon/format/orc/orc_memory_pool.h" #include "paimon/format/orc/orc_metrics.h" #include "paimon/format/orc/orc_output_stream_impl.h" +#include "paimon/format/orc/orc_reader_builder.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/predicate/predicate_builder.h" #include "paimon/testing/utils/read_result_collector.h" @@ -48,6 +49,12 @@ namespace paimon::orc::test { +TEST(OrcReaderBuilderTest, RejectsNullMemoryPool) { + OrcReaderBuilder builder(/*options=*/{}, /*batch_size=*/10); + builder.WithMemoryPool(nullptr); + ASSERT_NOK_WITH_MSG(builder.Build(nullptr), "ORC reader memory pool is nullptr"); +} + std::string SerializeSchemaToString(const std::shared_ptr& schema) { std::shared_ptr serialized = arrow::ipc::SerializeSchema(*schema).ValueOrDie(); return std::string(reinterpret_cast(serialized->data()), diff --git a/src/paimon/format/orc/orc_format_defs.h b/src/paimon/format/orc/orc_format_defs.h index 46eb5e7d8..f4b2da696 100644 --- a/src/paimon/format/orc/orc_format_defs.h +++ b/src/paimon/format/orc/orc_format_defs.h @@ -32,7 +32,7 @@ namespace paimon::orc { struct OrcReadMemory { explicit OrcReadMemory(const std::shared_ptr& pool) - : arrow_pool(GetSharedArrowPool(pool)), orc_pool(std::make_shared(pool)) {} + : arrow_pool(GetArrowPool(pool)), orc_pool(std::make_shared(pool)) {} std::shared_ptr arrow_pool; std::shared_ptr<::orc::MemoryPool> orc_pool; diff --git a/src/paimon/format/orc/orc_reader_builder.h b/src/paimon/format/orc/orc_reader_builder.h index f1669acbc..603479c7d 100644 --- a/src/paimon/format/orc/orc_reader_builder.h +++ b/src/paimon/format/orc/orc_reader_builder.h @@ -38,12 +38,19 @@ class OrcReaderBuilder : public ReaderBuilder { options_(options) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { - read_memory_ = std::make_shared(pool); + if (pool == nullptr) { + read_memory_.reset(); + } else { + read_memory_ = std::make_shared(pool); + } return this; } Result> Build( const std::shared_ptr& path) const override { + if (read_memory_ == nullptr) { + return Status::Invalid("ORC reader memory pool is nullptr"); + } PAIMON_ASSIGN_OR_RAISE(uint64_t natural_read_size, OptionsUtils::GetValueFromMap( options_, ORC_NATURAL_READ_SIZE, DEFAULT_NATURAL_READ_SIZE)); diff --git a/src/paimon/format/orc/orc_reader_wrapper.cpp b/src/paimon/format/orc/orc_reader_wrapper.cpp index 4eed0f500..b094182a6 100644 --- a/src/paimon/format/orc/orc_reader_wrapper.cpp +++ b/src/paimon/format/orc/orc_reader_wrapper.cpp @@ -84,7 +84,7 @@ Result OrcReaderWrapper::Next() { std::shared_ptr array, OrcAdapter::AppendBatch(target_type_, orc_batch.get(), read_memory_->arrow_pool.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), read_memory_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), read_memory_)); next_row_ = GetRowNumber() + orc_batch->numElements; guard.Release(); } catch (const std::exception& e) { diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index 95f5cb75c..e3740b850 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -638,7 +638,7 @@ Result ParquetFileBatchReader::NextBatch() { std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), arrow_pool_)); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); read_rows_ += array->length(); read_batch_count_++; diff --git a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp index 6190a955c..7c5a8f470 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -76,6 +76,12 @@ class Predicate; namespace paimon::parquet::test { +TEST(ParquetReaderBuilderTest, RejectsNullMemoryPool) { + ParquetReaderBuilder builder(/*options=*/{}, /*batch_size=*/10); + builder.WithMemoryPool(nullptr); + ASSERT_NOK_WITH_MSG(builder.Build(nullptr), "Parquet reader memory pool is nullptr"); +} + std::string SerializeSchemaToString(const std::shared_ptr& schema) { std::shared_ptr serialized = arrow::ipc::SerializeSchema(*schema).ValueOrDie(); return std::string(reinterpret_cast(serialized->data()), diff --git a/src/paimon/format/parquet/parquet_reader_builder.h b/src/paimon/format/parquet/parquet_reader_builder.h index 394991369..e4d2ddb43 100644 --- a/src/paimon/format/parquet/parquet_reader_builder.h +++ b/src/paimon/format/parquet/parquet_reader_builder.h @@ -51,12 +51,16 @@ class ParquetReaderBuilder : public ReaderBuilder { ParquetReaderBuilder(const std::map& options, int32_t batch_size) : batch_size_(batch_size), pool_(GetDefaultPool()), - arrow_pool_(GetSharedArrowPool(pool_)), + arrow_pool_(GetArrowPool(pool_)), options_(options) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; - arrow_pool_ = GetSharedArrowPool(pool); + if (pool == nullptr) { + arrow_pool_.reset(); + } else { + arrow_pool_ = GetArrowPool(pool); + } return this; } @@ -72,6 +76,9 @@ class ParquetReaderBuilder : public ReaderBuilder { Result> Build( const std::shared_ptr& path) const override { + if (pool_ == nullptr) { + return Status::Invalid("Parquet reader memory pool is nullptr"); + } try { PAIMON_ASSIGN_OR_RAISE(int64_t file_length, path->Length()); std::string file_uri; diff --git a/src/paimon/format/parquet/parquet_vector_io_test.cpp b/src/paimon/format/parquet/parquet_vector_io_test.cpp index 3e6ca2cb4..9c294f2a6 100644 --- a/src/paimon/format/parquet/parquet_vector_io_test.cpp +++ b/src/paimon/format/parquet/parquet_vector_io_test.cpp @@ -107,7 +107,7 @@ class ParquetVectorIoTest : public ::testing::Test { fs_->Create(file_path, /*overwrite=*/false)); ::parquet::WriterProperties::Builder properties_builder; properties_builder.max_row_group_length(max_row_group_length); - std::shared_ptr arrow_pool = GetSharedArrowPool(pool_); + std::shared_ptr arrow_pool = GetArrowPool(pool_); ASSERT_OK_AND_ASSIGN( std::unique_ptr writer, ParquetFormatWriter::Create(out, arrow::schema(write_type->fields()), @@ -138,7 +138,7 @@ class ParquetVectorIoTest : public ::testing::Test { fs_->Create(file_path, /*overwrite=*/false)); auto arrow_out = std::make_shared(out); ::parquet::WriterProperties::Builder properties_builder; - std::shared_ptr arrow_pool = GetSharedArrowPool(pool_); + std::shared_ptr arrow_pool = GetArrowPool(pool_); std::shared_ptr<::parquet::ArrowWriterProperties> arrow_properties = ::parquet::ArrowWriterProperties::Builder().store_schema()->build(); arrow::Status status = ::parquet::arrow::WriteTable( @@ -152,7 +152,7 @@ class ParquetVectorIoTest : public ::testing::Test { std::shared_ptr* file_type_out) { ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path)); ASSERT_OK_AND_ASSIGN(int64_t length, in->Length()); - std::shared_ptr arrow_pool = GetSharedArrowPool(pool_); + std::shared_ptr arrow_pool = GetArrowPool(pool_); auto in_stream = std::make_shared(in, length, arrow_pool); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, @@ -175,7 +175,7 @@ class ParquetVectorIoTest : public ::testing::Test { std::unique_ptr* vector_reader_out) { ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path)); ASSERT_OK_AND_ASSIGN(int64_t length, in->Length()); - std::shared_ptr arrow_pool = GetSharedArrowPool(pool_); + std::shared_ptr arrow_pool = GetArrowPool(pool_); auto in_stream = std::make_shared(in, length, arrow_pool); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader,