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/reader/batch_reader.h
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ 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
/// avoid potential issues during conversion through the Arrow C Data Interface.
///
/// @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 @@ -52,6 +54,8 @@ 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
/// avoid potential issues during conversion through the Arrow C Data Interface.
///
/// @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
4 changes: 2 additions & 2 deletions src/paimon/common/reader/data_evolution_file_reader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -165,8 +165,8 @@ Result<std::shared_ptr<arrow::Array>> DataEvolutionFileReader::NextBatchForSingl
if (concat_array_vec.empty()) {
return std::shared_ptr<arrow::Array>();
}
if (concat_array_vec.size() == 1) {
// avoid data copy
if (concat_array_vec.size() == 1 && concat_array_vec[0]->offset() == 0) {
// Avoid data copy when the array is already normalized.
return concat_array_vec[0];
}
// TODO(xinyu.lxy) remove data copy for efficiency
Expand Down
39 changes: 39 additions & 0 deletions src/paimon/common/utils/arrow/arrow_utils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,30 @@

#include "arrow/array/array_base.h"
#include "arrow/array/array_nested.h"
#include "arrow/array/concatenate.h"
#include "arrow/util/compression.h"
#include "fmt/format.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/string_utils.h"

namespace paimon {

namespace {

bool HasNonZeroOffset(const std::shared_ptr<arrow::ArrayData>& data) {
if (data->offset != 0) {
return true;
}
for (const auto& child : data->child_data) {
if (HasNonZeroOffset(child)) {
return true;
}
}
return false;
}

} // namespace

const char* ArrowUtils::kArrowSchemaMetadataKey = "ARROW:schema";

Result<std::shared_ptr<arrow::Schema>> ArrowUtils::DataTypeToSchema(
Expand Down Expand Up @@ -167,6 +184,28 @@ Result<std::shared_ptr<arrow::StructArray>> ArrowUtils::RemoveFieldFromStructArr
return array;
}

Result<std::shared_ptr<arrow::RecordBatch>> ArrowUtils::NormalizeRecordBatchOffsets(
const std::shared_ptr<arrow::RecordBatch>& record_batch, arrow::MemoryPool* pool) {
arrow::ArrayVector normalized_columns;
for (int32_t i = 0; i < record_batch->num_columns(); ++i) {
const std::shared_ptr<arrow::Array>& column = record_batch->column(i);
if (!HasNonZeroOffset(column->data())) {
continue;
}
if (normalized_columns.empty()) {
normalized_columns = record_batch->columns();
}
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> normalized_column,
arrow::Concatenate({column}, pool));
normalized_columns[i] = std::move(normalized_column);
}
if (normalized_columns.empty()) {
return record_batch;
}
return arrow::RecordBatch::Make(record_batch->schema(), record_batch->num_rows(),
std::move(normalized_columns));
}

Result<arrow::Compression::type> ArrowUtils::GetCompressionType(const std::string& compression) {
std::string normalized = StringUtils::ToLowerCase(compression);
if (normalized.empty() || normalized == "none") {
Expand Down
3 changes: 3 additions & 0 deletions src/paimon/common/utils/arrow/arrow_utils.h
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,9 @@ class PAIMON_EXPORT ArrowUtils {
static Result<std::shared_ptr<arrow::StructArray>> RemoveFieldFromStructArray(
const std::shared_ptr<arrow::StructArray>& struct_array, const std::string& field_name);

static Result<std::shared_ptr<arrow::RecordBatch>> NormalizeRecordBatchOffsets(
const std::shared_ptr<arrow::RecordBatch>& record_batch, arrow::MemoryPool* pool);

static bool EqualsIgnoreNullable(const std::shared_ptr<arrow::DataType>& type,
const std::shared_ptr<arrow::DataType>& other_type);

Expand Down
49 changes: 49 additions & 0 deletions src/paimon/common/utils/arrow/arrow_utils_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -379,6 +379,54 @@ TEST(ArrowUtilsTest, TestRemoveFieldFromStructArraySuccess) {
ASSERT_TRUE(result->Equals(expected_struct_array));
}

TEST(ArrowUtilsTest, TestNormalizeRecordBatchOffsets) {
auto value_field = arrow::field("value", arrow::int32());
auto nested_field = arrow::field("nested", arrow::struct_({value_field}));
auto text_field = arrow::field("text", arrow::utf8());
auto clean_field = arrow::field("clean", arrow::boolean());
auto schema = arrow::schema({nested_field, text_field, clean_field});

std::shared_ptr<arrow::Array> values =
arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 2, 3, 4]").ValueOrDie();
std::shared_ptr<arrow::Array> sliced_values = values->Slice(1, 3);
std::shared_ptr<arrow::StructArray> nested_column =
arrow::StructArray::Make({sliced_values}, {value_field->name()}).ValueOrDie();
std::shared_ptr<arrow::Array> text =
arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "c", "d", "e"])")
.ValueOrDie();
std::shared_ptr<arrow::Array> sliced_text = text->Slice(1, 3);
std::shared_ptr<arrow::Array> clean_column =
arrow::ipc::internal::json::ArrayFromJSON(arrow::boolean(), "[true, false, true]")
.ValueOrDie();
std::shared_ptr<arrow::RecordBatch> record_batch = arrow::RecordBatch::Make(
schema, /*num_rows=*/3, {nested_column, sliced_text, clean_column});

ASSERT_EQ(nested_column->offset(), 0);
ASSERT_EQ(nested_column->field(0)->offset(), 1);
ASSERT_EQ(sliced_text->offset(), 1);
ASSERT_EQ(clean_column->offset(), 0);

ASSERT_OK_AND_ASSIGN(
std::shared_ptr<arrow::RecordBatch> normalized_batch,
ArrowUtils::NormalizeRecordBatchOffsets(record_batch, arrow::default_memory_pool()));
ASSERT_NE(normalized_batch.get(), record_batch.get());
ASSERT_TRUE(normalized_batch->Equals(*record_batch));
std::shared_ptr<arrow::StructArray> normalized_nested =
std::static_pointer_cast<arrow::StructArray>(normalized_batch->column(0));
ASSERT_EQ(normalized_nested->offset(), 0);
ASSERT_EQ(normalized_nested->field(0)->offset(), 0);
ASSERT_EQ(normalized_batch->column(1)->offset(), 0);
ASSERT_EQ(normalized_batch->column(2)->offset(), 0);
ASSERT_NE(normalized_batch->column_data(0).get(), record_batch->column_data(0).get());
ASSERT_NE(normalized_batch->column_data(1).get(), record_batch->column_data(1).get());
ASSERT_EQ(normalized_batch->column_data(2).get(), record_batch->column_data(2).get());

ASSERT_OK_AND_ASSIGN(
std::shared_ptr<arrow::RecordBatch> unchanged_batch,
ArrowUtils::NormalizeRecordBatchOffsets(normalized_batch, arrow::default_memory_pool()));
ASSERT_EQ(unchanged_batch.get(), normalized_batch.get());
}

TEST(ArrowUtilsTest, TestEqualsIgnoreNullable) {
{
// test simple
Expand Down Expand Up @@ -481,6 +529,7 @@ TEST(ArrowUtilsTest, TestGetCompressionType) {
ASSERT_EQ(type, arrow::Compression::GZIP);
}
{
// test invalid codec
ASSERT_NOK(ArrowUtils::GetCompressionType("invalid_codec"));
}
}
Expand Down
5 changes: 5 additions & 0 deletions src/paimon/format/parquet/file_reader_wrapper.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
#include "arrow/record_batch.h"
#include "arrow/util/range.h"
#include "fmt/format.h"
#include "paimon/common/utils/arrow/arrow_utils.h"
#include "paimon/format/parquet/column_index_filter.h"
#include "paimon/format/parquet/page_filtered_row_group_reader.h"
#include "paimon/format/parquet/parquet_format_defs.h"
Expand Down Expand Up @@ -273,6 +274,10 @@ Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextFullyMatched(
if (!record_batch) {
return std::shared_ptr<arrow::RecordBatch>();
}
// Large binary columns (exceed 2GB) may split at different row boundaries. TableBatchReader
// aligns their chunks by slicing columns, which may leave non-zero child offsets.
PAIMON_ASSIGN_OR_RAISE(record_batch,
ArrowUtils::NormalizeRecordBatchOffsets(record_batch, pool_.get()));

int32_t rg_id = target_row_groups_[current_row_group_idx_].GetRowGroupIndex();
uint64_t rg_end = all_row_group_ranges_[rg_id].second;
Expand Down
32 changes: 25 additions & 7 deletions src/paimon/format/parquet/page_filtered_row_group_reader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
#include "arrow/table.h"
#include "arrow/util/future.h"
#include "fmt/format.h"
#include "paimon/common/utils/arrow/arrow_utils.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "parquet/arrow/reader_internal.h"
#include "parquet/metadata.h"
Expand All @@ -36,12 +37,13 @@ namespace paimon::parquet {
namespace {

/// Wraps an arrow::Table + TableBatchReader as a RecordBatchReader so the caller can
/// stream zero-copy-sliced batches without deep-copying multi-chunk columns. The Table
/// is held to keep its ChunkedArrays alive for the inner TableBatchReader.
/// stream batches while ensuring every returned array offset is zero. The Table is held
/// to keep its ChunkedArrays alive for the inner TableBatchReader.
class TableRecordBatchReader : public arrow::RecordBatchReader {
public:
TableRecordBatchReader(std::shared_ptr<arrow::Table> table, int64_t chunksize)
: table_(std::move(table)), inner_(*table_) {
TableRecordBatchReader(std::shared_ptr<arrow::Table> table, int64_t chunksize,
std::shared_ptr<arrow::MemoryPool> pool)
: table_(std::move(table)), inner_(*table_), pool_(std::move(pool)) {
inner_.set_chunksize(chunksize);
}

Expand All @@ -50,12 +52,26 @@ class TableRecordBatchReader : public arrow::RecordBatchReader {
}

arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch>* out) override {
return inner_.ReadNext(out);
ARROW_RETURN_NOT_OK(inner_.ReadNext(out));
if (!*out) {
return arrow::Status::OK();
}

// Page filtering may produce columns with different chunk boundaries. TableBatchReader
// aligns them by slicing columns, which may leave non-zero child offsets.
Result<std::shared_ptr<arrow::RecordBatch>> normalized_result =
ArrowUtils::NormalizeRecordBatchOffsets(*out, pool_.get());
if (!normalized_result.ok()) {
return ToArrowStatus(normalized_result.status());
}
*out = std::move(normalized_result).value();
return arrow::Status::OK();
}

private:
std::shared_ptr<arrow::Table> table_;
arrow::TableBatchReader inner_;
std::shared_ptr<arrow::MemoryPool> pool_;
};

} // namespace
Expand Down Expand Up @@ -240,7 +256,8 @@ Result<std::unique_ptr<arrow::RecordBatchReader>> PageFilteredRowGroupReader::Re
if (row_ranges.IsEmpty()) {
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Table> empty_table,
arrow::Table::MakeEmpty(arrow_schema, pool.get()));
return std::make_unique<TableRecordBatchReader>(std::move(empty_table), max_chunksize);
return std::make_unique<TableRecordBatchReader>(std::move(empty_table), max_chunksize,
pool);
}

int64_t expected_rows = row_ranges.RowCount();
Expand Down Expand Up @@ -283,7 +300,8 @@ Result<std::unique_ptr<arrow::RecordBatchReader>> PageFilteredRowGroupReader::Re
}

auto table = arrow::Table::Make(arrow_schema, std::move(columns), expected_rows);
return std::make_unique<TableRecordBatchReader>(std::move(table), max_chunksize);
return std::make_unique<TableRecordBatchReader>(std::move(table), max_chunksize,
std::move(pool));
}

std::vector<::arrow::io::ReadRange> PageFilteredRowGroupReader::ComputePageRanges(
Expand Down
24 changes: 20 additions & 4 deletions src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -659,14 +659,13 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesMultiplePages) {
ASSERT_LT(ranges[0].offset, ranges[1].offset);
}

/// Test: variable-length columns are streamed across multiple zero-copy-sliced
/// Test: variable-length columns are streamed across multiple offset-normalized
/// RecordBatches when batch_size is smaller than the matched row count, instead of
/// being concatenated into a single RecordBatch via CombineChunks.
///
/// This verifies the alignment with Arrow's standard TableBatchReader path:
/// multi-chunk binary/string columns split along chunk + batch_size boundaries,
/// with no deep copy. Asserts both correctness (total rows + full content order) and
/// the multi-batch shape (more than one chunk in the collected ChunkedArray).
/// multi-chunk binary/string columns split along chunk + batch_size boundaries. It
/// asserts correctness and the multi-batch shape.
TEST_F(PageFilteredRowGroupReaderTest, StringColumnMultiBatchStreaming) {
std::string file_name = dir_->Str() + "/string_multi_batch.parquet";

Expand Down Expand Up @@ -719,6 +718,23 @@ TEST_F(PageFilteredRowGroupReaderTest, StringColumnMultiBatchStreaming) {
ASSERT_EQ(40, seen);
}

TEST_F(PageFilteredRowGroupReaderTest, NormalizesSlicedBatchOffsets) {
std::string file_name = dir_->Str() + "/normalized_sliced_offsets.parquet";
std::shared_ptr<arrow::StructArray> data = MakeSequentialIntData(60);
WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/60);

std::shared_ptr<arrow::Schema> read_schema =
arrow::schema({arrow::field("val", arrow::int32())});
std::shared_ptr<Predicate> predicate = PredicateBuilder::GreaterOrEqual(
/*field_index=*/0, /*field_name=*/"val", FieldType::INT, Literal(20));

std::shared_ptr<arrow::ChunkedArray> result;
ReadWithPredicateImpl(file_name, read_schema, predicate, &result, /*batch_size=*/7);
ASSERT_TRUE(result);
ASSERT_EQ(40, result->length());
ASSERT_GT(result->num_chunks(), 1);
}

/// Test: end-to-end page-filtered read produces correct results when using page-level PreBuffer.
///
/// This exercises the full path: ComputePageRanges → PreBufferRanges → CachedInputStream →
Expand Down
19 changes: 19 additions & 0 deletions src/paimon/testing/utils/read_result_collector.h
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,7 @@ class ReadResultCollector {
if (BatchReader::IsEofBatch(batch_with_bitmap)) {
return std::shared_ptr<arrow::Array>();
}
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),
Expand All @@ -192,12 +193,30 @@ class ReadResultCollector {
if (BatchReader::IsEofBatch(batch)) {
return std::shared_ptr<arrow::Array>();
}
PAIMON_RETURN_NOT_OK(CheckBatchOffset(batch));
}
auto& [c_array, c_schema] = batch;
assert(c_array->length > 0);
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(auto result_array,
arrow::ImportArray(c_array.get(), c_schema.get()));
return result_array;
}

static Status CheckBatchOffset(const BatchReader::ReadBatch& batch) {
assert(!BatchReader::IsEofBatch(batch));
return CheckArrayOffset(batch.first.get());
}

static Status CheckArrayOffset(const ArrowArray* array) {
assert(array);
if (array->offset != 0) {
return Status::Invalid("BatchReader returned an array with non-zero offset " +
std::to_string(array->offset));
}
for (int64_t i = 0; i < array->n_children; i++) {
PAIMON_RETURN_NOT_OK(CheckArrayOffset(array->children[i]));
}
return Status::OK();
}
};
} // namespace paimon::test
Loading