Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions include/paimon/format/reader_builder.h
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::unique_ptr<FileBatchReader>> Build(
const std::shared_ptr<InputStream>& path) const = 0;
};
Expand Down
12 changes: 12 additions & 0 deletions include/paimon/reader/batch_reader.h
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand All @@ -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
Expand Down
14 changes: 1 addition & 13 deletions include/paimon/table/source/table_read.h
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
#include <vector>

#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"
Expand All @@ -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.
Expand Down Expand Up @@ -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<std::unique_ptr<BatchReader>> CreateReader(
const std::vector<std::shared_ptr<Split>>& splits);
const std::vector<std::shared_ptr<Split>>& splits) = 0;

/// Creates a `BatchReader` instance for a single split.
///
Expand All @@ -76,15 +74,5 @@ class PAIMON_EXPORT TableRead {
/// Implementations may override this to provide a more efficient count path.
virtual Result<std::unique_ptr<CountReader>> CreateCountReader(
const std::vector<std::shared_ptr<Split>>& splits);

protected:
explicit TableRead(const std::shared_ptr<MemoryPool>& memory_pool);

std::shared_ptr<MemoryPool> GetMemoryPool() const {
return pool_;
}

private:
std::shared_ptr<MemoryPool> pool_;
};
} // namespace paimon
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -137,7 +138,7 @@ class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test {
field_read_plans.emplace(field->name(), std::move(field_read_plan));
}
return std::make_unique<ShreddingFileReader>(std::move(reader), std::move(field_read_plans),
pool_);
GetSharedArrowPool(pool_));
}

Result<std::unique_ptr<ShreddingFileReader>> CreateReader(
Expand Down Expand Up @@ -246,7 +247,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
Expand All @@ -268,7 +269,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
Expand Down Expand Up @@ -301,14 +302,14 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjection)
selected_field, TagsMeta()));
std::map<std::string, std::shared_ptr<ShreddingColumnReadPlan>> contexts;
contexts.emplace("tags", std::move(field_read_plan));
auto reader =
std::make_unique<ShreddingFileReader>(std::move(mock_reader), std::move(contexts), pool_);
auto reader = std::make_unique<ShreddingFileReader>(std::move(mock_reader), std::move(contexts),
GetSharedArrowPool(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<arrow::ChunkedArray> expected;
Expand Down Expand Up @@ -441,14 +442,14 @@ TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjectionFr
file_schema->field(1), selected_field));
std::map<std::string, std::shared_ptr<ShreddingColumnReadPlan>> contexts;
contexts.emplace("tags", std::move(field_read_plan));
auto reader =
std::make_unique<ShreddingFileReader>(std::move(mock_reader), std::move(contexts), pool_);
auto reader = std::make_unique<ShreddingFileReader>(std::move(mock_reader), std::move(contexts),
GetSharedArrowPool(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<arrow::ChunkedArray> expected;
Expand Down Expand Up @@ -494,7 +495,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
Expand All @@ -514,7 +515,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
Expand Down Expand Up @@ -551,7 +552,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
Expand Down Expand Up @@ -590,7 +591,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
Expand Down Expand Up @@ -618,7 +619,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");
}

Expand All @@ -635,7 +636,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");
}

Expand Down Expand Up @@ -680,7 +681,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
Expand Down Expand Up @@ -741,7 +742,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
arrow::struct_(logical_schema->fields()), {R"([
Expand Down Expand Up @@ -801,7 +802,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
arrow::struct_(logical_schema->fields()), {R"([
Expand Down Expand Up @@ -852,7 +853,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<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
Expand Down
5 changes: 3 additions & 2 deletions src/paimon/common/data/shredding/shredding_file_reader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ namespace paimon {
ShreddingFileReader::ShreddingFileReader(
std::unique_ptr<FileBatchReader>&& reader,
std::map<std::string, std::shared_ptr<ShreddingColumnReadPlan>>&& plans,
const std::shared_ptr<MemoryPool>& pool)
: arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)), plans_(std::move(plans)) {}
const std::shared_ptr<arrow::MemoryPool>& arrow_pool)
: arrow_pool_(arrow_pool), reader_(std::move(reader)), plans_(std::move(plans)) {}

Result<std::unique_ptr<::ArrowSchema>> ShreddingFileReader::GetFileSchema() const {
return reader_->GetFileSchema();
Expand Down Expand Up @@ -109,6 +109,7 @@ Result<BatchReader::ReadBatchWithBitmap> ShreddingFileReader::NextBatchWithBitma
auto new_c_schema = std::make_unique<ArrowSchema>();
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;
}
Expand Down
2 changes: 1 addition & 1 deletion src/paimon/common/data/shredding/shredding_file_reader.h
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ class ShreddingFileReader : public FileBatchReader {
public:
ShreddingFileReader(std::unique_ptr<FileBatchReader>&& reader,
std::map<std::string, std::shared_ptr<ShreddingColumnReadPlan>>&& plans,
const std::shared_ptr<MemoryPool>& pool);
const std::shared_ptr<arrow::MemoryPool>& arrow_pool);

Result<std::unique_ptr<::ArrowSchema>> GetFileSchema() const override;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -60,10 +61,11 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test,
}
void TearDown() override {}

void CheckResult(BatchReader* apply_bitmap_batch_reader,
void CheckResult(std::unique_ptr<BatchReader> apply_bitmap_batch_reader,
const std::shared_ptr<arrow::ChunkedArray>& expected_chunk_array) {
ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::ChunkedArray> result_chunk_array,
ReadResultCollector::CollectResult(apply_bitmap_batch_reader));
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<arrow::ChunkedArray> 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());
Expand Down Expand Up @@ -98,22 +100,22 @@ 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_, GetSharedArrowPool(pool_)));
} else {
file_batch_reader =
std::make_unique<MockFileBatchReader>(data, target_type_, batch_size);
}
auto apply_bitmap_batch_reader = std::make_unique<ApplyBitmapIndexBatchReader>(
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<arrow::Array> expect_array =
arrow::StructArray::Make({expected}, target_type_->fields()).ValueOrDie();
auto expected_chunk_array = std::make_shared<arrow::ChunkedArray>(expect_array);
CheckResult(apply_bitmap_batch_reader.get(), expected_chunk_array);
CheckResult(std::move(apply_bitmap_batch_reader), expected_chunk_array);
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,13 +38,16 @@
namespace paimon {
CompleteIndexScoreBatchReader::CompleteIndexScoreBatchReader(
std::unique_ptr<BatchReader>&& reader, const std::vector<float>& scores,
const std::shared_ptr<MemoryPool>& pool)
: arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)), scores_(scores) {}
const std::shared_ptr<arrow::MemoryPool>& arrow_pool)
: arrow_pool_(arrow_pool), reader_(std::move(reader)), scores_(scores) {}

Result<BatchReader::ReadBatch> 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) {
Expand Down Expand Up @@ -106,6 +109,7 @@ Result<BatchReader::ReadBatchWithBitmap> 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
Loading
Loading