diff --git a/include/paimon/reader/batch_reader.h b/include/paimon/reader/batch_reader.h index 7378512f3..df5a99a27 100644 --- a/include/paimon/reader/batch_reader.h +++ b/include/paimon/reader/batch_reader.h @@ -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` @@ -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 diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index f33aea3af..f333b968a 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -18,6 +18,7 @@ #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" @@ -25,6 +26,22 @@ namespace paimon { +namespace { + +bool HasNonZeroOffset(const std::shared_ptr& 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> ArrowUtils::DataTypeToSchema( @@ -167,6 +184,28 @@ Result> ArrowUtils::RemoveFieldFromStructArr return array; } +Result> ArrowUtils::NormalizeRecordBatchOffsets( + const std::shared_ptr& record_batch, arrow::MemoryPool* pool) { + arrow::ArrayVector normalized_columns; + for (int32_t i = 0; i < record_batch->num_columns(); ++i) { + const std::shared_ptr& 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 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 ArrowUtils::GetCompressionType(const std::string& compression) { std::string normalized = StringUtils::ToLowerCase(compression); if (normalized.empty() || normalized == "none") { diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h index baf51d441..6e489023f 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -48,6 +48,9 @@ class PAIMON_EXPORT ArrowUtils { static Result> RemoveFieldFromStructArray( const std::shared_ptr& struct_array, const std::string& field_name); + static Result> NormalizeRecordBatchOffsets( + const std::shared_ptr& record_batch, arrow::MemoryPool* pool); + static bool EqualsIgnoreNullable(const std::shared_ptr& type, const std::shared_ptr& other_type); diff --git a/src/paimon/common/utils/arrow/arrow_utils_test.cpp b/src/paimon/common/utils/arrow/arrow_utils_test.cpp index dfeb8c504..86b4c016c 100644 --- a/src/paimon/common/utils/arrow/arrow_utils_test.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils_test.cpp @@ -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 values = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 2, 3, 4]").ValueOrDie(); + std::shared_ptr sliced_values = values->Slice(1, 3); + std::shared_ptr nested_column = + arrow::StructArray::Make({sliced_values}, {value_field->name()}).ValueOrDie(); + std::shared_ptr text = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "c", "d", "e"])") + .ValueOrDie(); + std::shared_ptr sliced_text = text->Slice(1, 3); + std::shared_ptr clean_column = + arrow::ipc::internal::json::ArrayFromJSON(arrow::boolean(), "[true, false, true]") + .ValueOrDie(); + std::shared_ptr 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 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 normalized_nested = + std::static_pointer_cast(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 unchanged_batch, + ArrowUtils::NormalizeRecordBatchOffsets(normalized_batch, arrow::default_memory_pool())); + ASSERT_EQ(unchanged_batch.get(), normalized_batch.get()); +} + TEST(ArrowUtilsTest, TestEqualsIgnoreNullable) { { // test simple @@ -481,6 +529,7 @@ TEST(ArrowUtilsTest, TestGetCompressionType) { ASSERT_EQ(type, arrow::Compression::GZIP); } { + // test invalid codec ASSERT_NOK(ArrowUtils::GetCompressionType("invalid_codec")); } } diff --git a/src/paimon/format/parquet/file_reader_wrapper.cpp b/src/paimon/format/parquet/file_reader_wrapper.cpp index 090da9122..21cfa930b 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.cpp +++ b/src/paimon/format/parquet/file_reader_wrapper.cpp @@ -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" @@ -273,6 +274,10 @@ Result> FileReaderWrapper::NextFullyMatched( if (!record_batch) { return std::shared_ptr(); } + // 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; diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp index 20c5efb97..5c6fb8653 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp @@ -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" @@ -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 table, int64_t chunksize) - : table_(std::move(table)), inner_(*table_) { + TableRecordBatchReader(std::shared_ptr table, int64_t chunksize, + std::shared_ptr pool) + : table_(std::move(table)), inner_(*table_), pool_(std::move(pool)) { inner_.set_chunksize(chunksize); } @@ -50,12 +52,26 @@ class TableRecordBatchReader : public arrow::RecordBatchReader { } arrow::Status ReadNext(std::shared_ptr* 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> 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 table_; arrow::TableBatchReader inner_; + std::shared_ptr pool_; }; } // namespace @@ -240,7 +256,8 @@ Result> PageFilteredRowGroupReader::Re if (row_ranges.IsEmpty()) { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr empty_table, arrow::Table::MakeEmpty(arrow_schema, pool.get())); - return std::make_unique(std::move(empty_table), max_chunksize); + return std::make_unique(std::move(empty_table), max_chunksize, + pool); } int64_t expected_rows = row_ranges.RowCount(); @@ -283,7 +300,8 @@ Result> PageFilteredRowGroupReader::Re } auto table = arrow::Table::Make(arrow_schema, std::move(columns), expected_rows); - return std::make_unique(std::move(table), max_chunksize); + return std::make_unique(std::move(table), max_chunksize, + std::move(pool)); } std::vector<::arrow::io::ReadRange> PageFilteredRowGroupReader::ComputePageRanges( 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 46347f33f..361c8b199 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 @@ -661,14 +661,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"; @@ -721,6 +720,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 data = MakeSequentialIntData(60); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/60); + + std::shared_ptr read_schema = + arrow::schema({arrow::field("val", arrow::int32())}); + std::shared_ptr predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/0, /*field_name=*/"val", FieldType::INT, Literal(20)); + + std::shared_ptr 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 → diff --git a/src/paimon/testing/utils/read_result_collector.h b/src/paimon/testing/utils/read_result_collector.h index 986176119..784f03155 100644 --- a/src/paimon/testing/utils/read_result_collector.h +++ b/src/paimon/testing/utils/read_result_collector.h @@ -180,6 +180,7 @@ class ReadResultCollector { if (BatchReader::IsEofBatch(batch_with_bitmap)) { return std::shared_ptr(); } + 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), @@ -192,6 +193,7 @@ class ReadResultCollector { if (BatchReader::IsEofBatch(batch)) { return std::shared_ptr(); } + PAIMON_RETURN_NOT_OK(CheckBatchOffset(batch)); } auto& [c_array, c_schema] = batch; assert(c_array->length > 0); @@ -199,5 +201,22 @@ class ReadResultCollector { 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