This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-cpp.git
The following commit(s) were added to refs/heads/main by this push:
new f75a464 perf(reader): rebase Arrow offsets by slicing buffers
(zero-copy) instead of copying values (#188)
f75a464 is described below
commit f75a464303c8598a9be04fa0c464c5a69626db64
Author: Zhou Hongfeng <[email protected]>
AuthorDate: Wed Aug 12 13:56:09 2026 +0800
perf(reader): rebase Arrow offsets by slicing buffers (zero-copy) instead
of copying values (#188)
* feat: support zero-copy slice for RecordBatch
* test: add test cases
* refractor: extract RebaseBoolean function
* style: clang-tiyd
* add test for large_list
* stlye: add TODO
* style: pre-commit
* style: pre-commit
* style: change variable type
* test: add test case for timestamp
* test: add test cases for nested columns with offsets is zero while
children type is non-zero
---
src/paimon/common/utils/arrow/arrow_utils.cpp | 212 +++++++++++++++-
src/paimon/common/utils/arrow/arrow_utils.h | 3 +
src/paimon/common/utils/arrow/arrow_utils_test.cpp | 274 +++++++++++++++++++++
.../parquet/page_filtered_row_group_reader.cpp | 7 +-
4 files changed, 492 insertions(+), 4 deletions(-)
diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp
b/src/paimon/common/utils/arrow/arrow_utils.cpp
index de6aedb..04ce721 100644
--- a/src/paimon/common/utils/arrow/arrow_utils.cpp
+++ b/src/paimon/common/utils/arrow/arrow_utils.cpp
@@ -22,6 +22,12 @@
#include "arrow/array/array_base.h"
#include "arrow/array/array_nested.h"
#include "arrow/array/concatenate.h"
+#include "arrow/array/util.h"
+#include "arrow/buffer.h"
+#include "arrow/type_traits.h"
+#include "arrow/util/bit_util.h"
+#include "arrow/util/bitmap_ops.h"
+#include "arrow/util/checked_cast.h"
#include "arrow/util/compression.h"
#include "fmt/format.h"
#include "paimon/common/utils/arrow/status_utils.h"
@@ -43,6 +49,206 @@ bool HasNonZeroOffset(const
std::shared_ptr<arrow::ArrayData>& data) {
return false;
}
+Result<std::shared_ptr<arrow::ArrayData>> CopyToZeroOffset(
+ const std::shared_ptr<arrow::ArrayData>& data, arrow::MemoryPool* pool) {
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> copied,
+
arrow::Concatenate({arrow::MakeArray(data)}, pool));
+ return copied->data();
+}
+
+Result<std::shared_ptr<arrow::Buffer>> RebaseBitmap(const arrow::ArrayData&
data,
+ const
std::shared_ptr<arrow::Buffer>& bitmap,
+ arrow::MemoryPool* pool) {
+ if (data.offset % 8 == 0) {
+ return arrow::SliceBuffer(bitmap, data.offset / 8,
+ arrow::bit_util::BytesForBits(data.length));
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::Buffer> copied,
+ arrow::internal::CopyBitmap(pool, bitmap->data(), data.offset,
data.length));
+ return copied;
+}
+
+Result<std::shared_ptr<arrow::Buffer>> RebaseValidityBitmap(const
arrow::ArrayData& data,
+ arrow::MemoryPool*
pool) {
+ const std::shared_ptr<arrow::Buffer>& bitmap = data.buffers[0];
+ if (bitmap == nullptr || data.null_count.load() == 0) {
+ return std::shared_ptr<arrow::Buffer>();
+ }
+ return RebaseBitmap(data, bitmap, pool);
+}
+
+struct RebasedOffsets {
+ std::shared_ptr<arrow::Buffer> buffer;
+ int64_t first_value = 0;
+ int64_t last_value = 0;
+};
+
+template <typename OffsetType>
+Result<RebasedOffsets> RebaseOffsets(const arrow::ArrayData& data,
arrow::MemoryPool* pool) {
+ const auto* offsets = data.GetValues<OffsetType>(1);
+ const OffsetType base = offsets[0];
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::unique_ptr<arrow::Buffer> buffer,
+ arrow::AllocateBuffer((data.length + 1) *
static_cast<int64_t>(sizeof(OffsetType)), pool));
+ auto* rebased = reinterpret_cast<OffsetType*>(buffer->mutable_data());
+ for (int64_t i = 0; i <= data.length; i++) {
+ rebased[i] = offsets[i] - base;
+ }
+ return RebasedOffsets{std::shared_ptr<arrow::Buffer>(std::move(buffer)),
base,
+ offsets[data.length]};
+}
+
+Result<std::shared_ptr<arrow::ArrayData>> RebaseToZeroOffset(
+ const std::shared_ptr<arrow::ArrayData>& data, arrow::MemoryPool* pool);
+
+/// Rebases a boolean array, whose values are a bitmap rather than byte
addressable.
+Result<std::shared_ptr<arrow::ArrayData>> RebaseBoolean(
+ const std::shared_ptr<arrow::ArrayData>& data, arrow::MemoryPool* pool) {
+ if (data->buffers.size() != 2 || data->buffers[1] == nullptr) {
+ return CopyToZeroOffset(data, pool);
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> validity,
+ RebaseValidityBitmap(*data, pool));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> values,
+ RebaseBitmap(*data, data->buffers[1], pool));
+ std::shared_ptr<arrow::ArrayData> rebased =
+ arrow::ArrayData::Make(data->type, data->length,
data->null_count.load(), /*offset=*/0);
+ rebased->buffers = {std::move(validity), std::move(values)};
+ return rebased;
+}
+
+/// Rebases the {validity, offsets, values} layout of binary-like arrays.
+template <typename OffsetType>
+Result<std::shared_ptr<arrow::ArrayData>> RebaseBinaryLike(
+ const std::shared_ptr<arrow::ArrayData>& data, arrow::MemoryPool* pool) {
+ if (data->buffers.size() != 3 || data->buffers[1] == nullptr ||
data->buffers[2] == nullptr) {
+ return CopyToZeroOffset(data, pool);
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> validity,
+ RebaseValidityBitmap(*data, pool));
+ PAIMON_ASSIGN_OR_RAISE(RebasedOffsets offsets,
RebaseOffsets<OffsetType>(*data, pool));
+ std::shared_ptr<arrow::ArrayData> rebased =
+ arrow::ArrayData::Make(data->type, data->length,
data->null_count.load(), /*offset=*/0);
+ rebased->buffers = {std::move(validity), std::move(offsets.buffer),
+ arrow::SliceBuffer(data->buffers[2],
offsets.first_value,
+ offsets.last_value -
offsets.first_value)};
+ return rebased;
+}
+
+/// Rebases the {validity, offsets} plus single child layout of list, large
list and map arrays.
+template <typename OffsetType>
+Result<std::shared_ptr<arrow::ArrayData>> RebaseListLike(
+ const std::shared_ptr<arrow::ArrayData>& data, arrow::MemoryPool* pool) {
+ if (data->buffers.size() != 2 || data->buffers[1] == nullptr ||
data->child_data.size() != 1) {
+ return CopyToZeroOffset(data, pool);
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> validity,
+ RebaseValidityBitmap(*data, pool));
+ PAIMON_ASSIGN_OR_RAISE(RebasedOffsets offsets,
RebaseOffsets<OffsetType>(*data, pool));
+ // A contiguous slice of the parent always spans a contiguous range of the
child.
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::ArrayData> child_slice,
+ data->child_data[0]->SliceSafe(offsets.first_value,
+ offsets.last_value -
offsets.first_value));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ArrayData> child,
+ RebaseToZeroOffset(child_slice, pool));
+ std::shared_ptr<arrow::ArrayData> rebased =
+ arrow::ArrayData::Make(data->type, data->length,
data->null_count.load(), /*offset=*/0);
+ rebased->buffers = {std::move(validity), std::move(offsets.buffer)};
+ rebased->child_data = {std::move(child)};
+ return rebased;
+}
+
+/// Rebases a struct array, whose slices keep full length children.
+Result<std::shared_ptr<arrow::ArrayData>> RebaseStruct(
+ const std::shared_ptr<arrow::ArrayData>& data, arrow::MemoryPool* pool) {
+ if (data->buffers.empty()) {
+ return CopyToZeroOffset(data, pool);
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> validity,
+ RebaseValidityBitmap(*data, pool));
+ std::shared_ptr<arrow::ArrayData> rebased =
+ arrow::ArrayData::Make(data->type, data->length,
data->null_count.load(), /*offset=*/0);
+ rebased->buffers = {std::move(validity)};
+ rebased->child_data.reserve(data->child_data.size());
+ for (const auto& child : data->child_data) {
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::ArrayData>
child_slice,
+ child->SliceSafe(data->offset,
data->length));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ArrayData> rebased_child,
+ RebaseToZeroOffset(child_slice, pool));
+ rebased->child_data.push_back(std::move(rebased_child));
+ }
+ return rebased;
+}
+
+/// Slices the single value buffer of a fixed width array. Returns nullptr
when the layout is not
+/// a plain byte addressable fixed width one.
+Result<std::shared_ptr<arrow::ArrayData>> RebaseFixedWidth(
+ const std::shared_ptr<arrow::ArrayData>& data, arrow::MemoryPool* pool) {
+ // arrow::is_fixed_width() also covers dictionary types, whose dictionary
is not in child_data.
+ if (!arrow::is_fixed_width(data->type->id()) || data->buffers.size() != 2
||
+ data->buffers[1] == nullptr || !data->child_data.empty() ||
data->dictionary != nullptr) {
+ return std::shared_ptr<arrow::ArrayData>();
+ }
+ const int32_t bit_width =
+ arrow::internal::checked_cast<const
arrow::FixedWidthType&>(*data->type).bit_width();
+ if (bit_width <= 0 || bit_width % 8 != 0) {
+ return std::shared_ptr<arrow::ArrayData>();
+ }
+ const int64_t byte_width = bit_width / 8;
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> validity,
+ RebaseValidityBitmap(*data, pool));
+ std::shared_ptr<arrow::ArrayData> rebased =
+ arrow::ArrayData::Make(data->type, data->length,
data->null_count.load(), /*offset=*/0);
+ rebased->buffers = {
+ std::move(validity),
+ arrow::SliceBuffer(data->buffers[1], data->offset * byte_width,
data->length * byte_width)};
+ return rebased;
+}
+
+/// Returns an ArrayData describing the same rows as `data` with a zero offset
at every level.
+/// Buffers are sliced rather than copied wherever the layout allows it, so
the cost is
+/// proportional to the number of rows instead of the number of value bytes.
+Result<std::shared_ptr<arrow::ArrayData>> RebaseToZeroOffset(
+ const std::shared_ptr<arrow::ArrayData>& data, arrow::MemoryPool* pool) {
+ if (!HasNonZeroOffset(data)) {
+ return data;
+ }
+ // An empty array may not carry the buffers the layouts below slice.
+ if (data->length == 0 || data->buffers.empty()) {
+ return CopyToZeroOffset(data, pool);
+ }
+
+ switch (data->type->id()) {
+ case arrow::Type::BOOL:
+ return RebaseBoolean(data, pool);
+ case arrow::Type::STRING:
+ case arrow::Type::BINARY:
+ return RebaseBinaryLike<int32_t>(data, pool);
+ case arrow::Type::LARGE_STRING:
+ case arrow::Type::LARGE_BINARY:
+ return RebaseBinaryLike<int64_t>(data, pool);
+ case arrow::Type::LIST:
+ case arrow::Type::MAP:
+ return RebaseListLike<int32_t>(data, pool);
+ case arrow::Type::LARGE_LIST:
+ return RebaseListLike<int64_t>(data, pool);
+ case arrow::Type::STRUCT:
+ return RebaseStruct(data, pool);
+ case arrow::Type::DICTIONARY:
+ return CopyToZeroOffset(data, pool);
+ default:
+ break;
+ }
+
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ArrayData> rebased,
RebaseFixedWidth(data, pool));
+ if (rebased != nullptr) {
+ return rebased;
+ }
+ return CopyToZeroOffset(data, pool);
+}
+
} // namespace
const char* ArrowUtils::kArrowSchemaMetadataKey = "ARROW:schema";
@@ -198,9 +404,9 @@ Result<std::shared_ptr<arrow::RecordBatch>>
ArrowUtils::NormalizeRecordBatchOffs
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);
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ArrayData>
normalized_data,
+ RebaseToZeroOffset(column->data(), pool));
+ normalized_columns[i] = arrow::MakeArray(normalized_data);
}
if (normalized_columns.empty()) {
return record_batch;
diff --git a/src/paimon/common/utils/arrow/arrow_utils.h
b/src/paimon/common/utils/arrow/arrow_utils.h
index 1d4b785..326b388 100644
--- a/src/paimon/common/utils/arrow/arrow_utils.h
+++ b/src/paimon/common/utils/arrow/arrow_utils.h
@@ -51,6 +51,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);
+ /// Returns a RecordBatch whose columns, including their nested children,
all have a zero
+ /// offset, as required by `BatchReader`. Offsets are rebased by slicing
buffers (zero copy)
+ /// wherever the layout allows it; only layouts that cannot be rebased
fall back to a full copy.
static Result<std::shared_ptr<arrow::RecordBatch>>
NormalizeRecordBatchOffsets(
const std::shared_ptr<arrow::RecordBatch>& record_batch,
arrow::MemoryPool* pool);
diff --git a/src/paimon/common/utils/arrow/arrow_utils_test.cpp
b/src/paimon/common/utils/arrow/arrow_utils_test.cpp
index a2f8d5b..032e8b4 100644
--- a/src/paimon/common/utils/arrow/arrow_utils_test.cpp
+++ b/src/paimon/common/utils/arrow/arrow_utils_test.cpp
@@ -430,6 +430,280 @@ TEST(ArrowUtilsTest, TestNormalizeRecordBatchOffsets) {
ASSERT_EQ(unchanged_batch.get(), normalized_batch.get());
}
+namespace {
+
+/// A buffer that rebasing must expose as a view into the source.
+// This struct tells where a ArrayData stores value.
+struct SharedBuffer {
+ std::vector<int32_t> child_path;
+ int buffer_index;
+};
+
+struct NormalizeCase {
+ std::shared_ptr<arrow::DataType> type;
+ std::string json;
+ /// The buffers holding the values of this layout, which rebasing must
never copy.
+ std::vector<SharedBuffer> value_buffers;
+};
+
+/// Ten values per case, so that the slices taken below stay in range.
+std::vector<NormalizeCase> NormalizeCases() {
+ auto int_field = arrow::field("a", arrow::int32());
+ auto text_field = arrow::field("b", arrow::utf8());
+ return {
+ {arrow::boolean(),
+ "[true, null, false, true, true, null, false, false, true, null]",
+ {{{}, 1}}},
+ {arrow::int8(), "[0, 1, null, 3, 4, 5, null, 7, 8, 9]", {{{}, 1}}},
+ {arrow::int32(), "[0, 1, null, 3, 4, 5, null, 7, 8, 9]", {{{}, 1}}},
+ {arrow::int64(), "[0, 1, null, 3, 4, 5, null, 7, 8, 9]", {{{}, 1}}},
+ {arrow::float64(), "[0.5, 1.5, null, 3.5, 4.5, 5.5, null, 7.5, 8.5,
9.5]", {{{}, 1}}},
+ {arrow::date32(), "[0, 1, null, 3, 4, 5, null, 7, 8, 9]", {{{}, 1}}},
+ {arrow::timestamp(arrow::TimeUnit::MICRO),
+ "[0, 1, null, 3, 4, 5, null, 7, 8, 9]",
+ {{{}, 1}}},
+ {arrow::timestamp(arrow::TimeUnit::MILLI, "UTC"),
+ "[0, 1, null, 3, 4, 5, null, 7, 8, 9]",
+ {{{}, 1}}},
+ {arrow::decimal128(10, 2),
+ R"(["1.23", null, "3.45", "6.78", "0.01", null, "9.99", "8.88",
"7.77", "6.66"])",
+ {{{}, 1}}},
+ // no nulls at all, so the validity bitmap is dropped rather than
rebased
+ {arrow::int32(), "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]", {{{}, 1}}},
+ {arrow::utf8(),
+ R"(["a", null, "ccc", "dddd", "", "ffffff", null, "h", "ii", "jjj"])",
+ {{{}, 2}}},
+ {arrow::binary(),
+ R"(["a", null, "ccc", "dddd", "", "ffffff", null, "h", "ii", "jjj"])",
+ {{{}, 2}}},
+ {arrow::large_binary(),
+ R"(["a", null, "ccc", "dddd", "", "ffffff", null, "h", "ii", "jjj"])",
+ {{{}, 2}}},
+ {arrow::list(arrow::int32()),
+ "[[1], null, [2, 3], [], [4, 5, 6], null, [7], [8, 9], [], [10]]",
+ {{{0}, 1}}},
+ {arrow::large_list(arrow::int32()),
+ "[[1], null, [2, 3], [], [4, 5, 6], null, [7], [8, 9], [], [10]]",
+ {{{0}, 1}}},
+ {arrow::list(arrow::utf8()),
+ R"([["a"], null, ["bb", "ccc"], [], ["d"], null, ["e", "f"], [],
["g"], ["h"]])",
+ {{{0}, 2}}},
+ {arrow::struct_({int_field, text_field}),
+ R"([{"a": 0, "b": "x"}, null, {"a": 2, "b": null}, {"a": null, "b":
"yyy"},
+ {"a": 4, "b": "z"}, {"a": 5, "b": ""}, null, {"a": 7, "b": "w"},
+ {"a": 8, "b": "vv"}, {"a": 9, "b": "u"}])",
+ {{{0}, 1}, {{1}, 2}}},
+ // a list of structs exercises two levels of rebasing at once
+ {arrow::list(arrow::struct_({int_field, text_field})),
+ R"([[{"a": 0, "b": "x"}], null, [{"a": 2, "b": "y"}, {"a": 3, "b":
null}], [],
+ [{"a": 4, "b": "z"}], null, [{"a": 6, "b": "w"}], [], [{"a": 8,
"b": "v"}],
+ [{"a": 9, "b": "u"}]])",
+ {{{0, 0}, 1}, {{0, 1}, 2}}},
+ {arrow::map(arrow::utf8(), arrow::int32()),
+ R"([[["k0", 0]], null, [["k1", 1], ["k2", null]], [], [["k3", 3]],
null,
+ [["k4", 4], ["k5", 5]], [], [["k6", 6]], [["k7", 7]]])",
+ {{{0, 0}, 2}, {{0, 1}, 1}}},
+ };
+}
+
+const arrow::ArrayData& ResolvePath(const arrow::ArrayData& data,
+ const std::vector<int32_t>& child_path) {
+ const arrow::ArrayData* node = &data;
+ for (int32_t child_index : child_path) {
+ node = node->child_data[child_index].get();
+ }
+ return *node;
+}
+
+void ExpectAllOffsetsZero(const arrow::ArrayData& data, const std::string&
path) {
+ ASSERT_EQ(data.offset, 0) << "non-zero offset at " << path;
+ for (size_t i = 0; i < data.child_data.size(); i++) {
+ ExpectAllOffsetsZero(*data.child_data[i], path + "/child" +
std::to_string(i));
+ }
+}
+
+/// A freshly allocated buffer cannot live inside a buffer that is still
alive, so containment
+/// proves that `rebased` references the source bytes instead of copying them.
+bool IsViewInto(const arrow::Buffer& rebased, const arrow::Buffer& source) {
+ return rebased.data() >= source.data() &&
+ rebased.data() + rebased.size() <= source.data() + source.size();
+}
+
+std::shared_ptr<arrow::RecordBatch> MakeSliceBatch(const
std::shared_ptr<arrow::Array>& array,
+ int64_t offset, int64_t
length) {
+ return arrow::RecordBatch::Make(arrow::schema({arrow::field("f",
array->type())}), length,
+ {array->Slice(offset, length)});
+}
+
+/// Checks that normalization keeps the same rows with every offset zeroed.
`array` is used as a
+/// single column batch, so its own offset is whatever the caller built it
with.
+void CheckNormalizedArray(const std::shared_ptr<arrow::Array>& array) {
+ SCOPED_TRACE("type=" + array->type()->ToString());
+ std::shared_ptr<arrow::RecordBatch> batch = arrow::RecordBatch::Make(
+ arrow::schema({arrow::field("f", array->type())}), array->length(),
{array});
+
+ ASSERT_OK_AND_ASSIGN(
+ std::shared_ptr<arrow::RecordBatch> normalized,
+ ArrowUtils::NormalizeRecordBatchOffsets(batch,
arrow::default_memory_pool()));
+ // A batch that needs normalization must not be returned unchanged.
+ ASSERT_NE(normalized.get(), batch.get());
+
+ arrow::Status validated = normalized->ValidateFull();
+ ASSERT_TRUE(validated.ok()) << validated.ToString();
+ ASSERT_TRUE(normalized->Equals(*batch))
+ << "expected " << batch->ToString() << " but got " <<
normalized->ToString();
+ ExpectAllOffsetsZero(*normalized->column_data(0), "f");
+}
+
+/// Slices `array` and checks that normalization keeps the same rows with
every offset zeroed.
+void CheckNormalizedSlice(const std::shared_ptr<arrow::Array>& array, int64_t
offset,
+ int64_t length) {
+ SCOPED_TRACE("type=" + array->type()->ToString() + " offset=" +
std::to_string(offset) +
+ " length=" + std::to_string(length));
+ std::shared_ptr<arrow::RecordBatch> batch = MakeSliceBatch(array, offset,
length);
+
+ ASSERT_OK_AND_ASSIGN(
+ std::shared_ptr<arrow::RecordBatch> normalized,
+ ArrowUtils::NormalizeRecordBatchOffsets(batch,
arrow::default_memory_pool()));
+
+ arrow::Status validated = normalized->ValidateFull();
+ ASSERT_TRUE(validated.ok()) << validated.ToString();
+ ASSERT_TRUE(normalized->Equals(*batch))
+ << "expected " << batch->ToString() << " but got " <<
normalized->ToString();
+ ExpectAllOffsetsZero(*normalized->column_data(0), "f");
+}
+
+} // namespace
+
+// Check the equality of normalized batches and original batches.
+TEST(ArrowUtilsTest, TestNormalizeRecordBatchOffsetsCoversSupportedTypes) {
+ for (const NormalizeCase& normalize_case : NormalizeCases()) {
+ SCOPED_TRACE("type=" + normalize_case.type->ToString());
+ std::shared_ptr<arrow::Array> array =
+ arrow::ipc::internal::json::ArrayFromJSON(normalize_case.type,
normalize_case.json)
+ .ValueOrDie();
+ ASSERT_EQ(array->length(), 10);
+ // offset 0 takes the no-op path, offsets 1/3/5/9 are not byte
aligned, offset 8 is
+ for (const auto& [offset, length] : std::vector<std::pair<int64_t,
int64_t>>{
+ {0, 10}, {1, 9}, {1, 3}, {3, 4}, {5, 5}, {8, 2}, {9, 1}, {2,
0}}) {
+ CheckNormalizedSlice(array, offset, length);
+ }
+ }
+}
+
+// Check the zero-copy property of value buffer rebasing.
+TEST(ArrowUtilsTest, TestNormalizeRecordBatchOffsetsSharesValueBuffers) {
+ for (const NormalizeCase& normalize_case : NormalizeCases()) {
+ SCOPED_TRACE("type=" + normalize_case.type->ToString());
+ std::shared_ptr<arrow::Array> array =
+ arrow::ipc::internal::json::ArrayFromJSON(normalize_case.type,
normalize_case.json)
+ .ValueOrDie();
+ // A byte aligned offset lets bitmaps be sliced too, so nothing has to
be copied here.
+ std::shared_ptr<arrow::RecordBatch> batch =
+ MakeSliceBatch(array, /*offset=*/8, /*length=*/2);
+
+ ASSERT_OK_AND_ASSIGN(
+ std::shared_ptr<arrow::RecordBatch> normalized,
+ ArrowUtils::NormalizeRecordBatchOffsets(batch,
arrow::default_memory_pool()));
+ ASSERT_NE(normalized.get(), batch.get());
+
+ for (const SharedBuffer& value_buffer : normalize_case.value_buffers) {
+ SCOPED_TRACE("buffer_index=" +
std::to_string(value_buffer.buffer_index));
+ const arrow::ArrayData& rebased =
+ ResolvePath(*normalized->column_data(0),
value_buffer.child_path);
+ const arrow::ArrayData& source = ResolvePath(*array->data(),
value_buffer.child_path);
+ ASSERT_TRUE(IsViewInto(*rebased.buffers[value_buffer.buffer_index],
+ *source.buffers[value_buffer.buffer_index]))
+ << "value buffer was copied instead of sliced";
+ }
+ }
+}
+
+// Check the situation that child offsets are non-zero while parent offset is
zero.
+TEST(ArrowUtilsTest,
TestNormalizeRecordBatchOffsetsRebasesNestedOffsetsUnderZeroParent) {
+ std::shared_ptr<arrow::Array> ints =
+ arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 2,
3, 4, 5, 6, 7]")
+ .ValueOrDie();
+ std::shared_ptr<arrow::Array> texts =
+ arrow::ipc::internal::json::ArrayFromJSON(
+ arrow::utf8(), R"(["a", "bb", null, "dddd", "e", "ff", "ggg",
"h"])")
+ .ValueOrDie();
+
+ {
+ // struct whose children are sliced: parent offset 0, both children
offset 2
+ std::shared_ptr<arrow::Array> array =
+ arrow::StructArray::Make({ints->Slice(2, 4), texts->Slice(2, 4)},
+ std::vector<std::string>{"a", "b"})
+ .ValueOrDie();
+ ASSERT_EQ(array->offset(), 0);
+ ASSERT_EQ(array->data()->child_data[0]->offset, 2);
+ ASSERT_EQ(array->data()->child_data[1]->offset, 2);
+ CheckNormalizedArray(array);
+ }
+ {
+ // only the innermost array is sliced, so detection has to walk two
levels down
+ std::shared_ptr<arrow::Array> inner =
+ arrow::StructArray::Make({ints->Slice(3, 4)},
std::vector<std::string>{"a"})
+ .ValueOrDie();
+ std::shared_ptr<arrow::Array> array =
+ arrow::StructArray::Make({inner},
std::vector<std::string>{"inner"}).ValueOrDie();
+ ASSERT_EQ(array->offset(), 0);
+ ASSERT_EQ(array->data()->child_data[0]->offset, 0);
+ ASSERT_EQ(array->data()->child_data[0]->child_data[0]->offset, 3);
+ CheckNormalizedArray(array);
+ }
+ {
+ // list built over sliced values: parent offset 0, values offset 2,
and the list offsets
+ // address the values relative to that slice
+ std::shared_ptr<arrow::Array> offsets =
+ arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1,
1, 3, 4]")
+ .ValueOrDie();
+ std::shared_ptr<arrow::Array> array =
+ arrow::ListArray::FromArrays(*offsets, *ints->Slice(2,
4)).ValueOrDie();
+ ASSERT_EQ(array->offset(), 0);
+ ASSERT_EQ(array->data()->child_data[0]->offset, 2);
+ CheckNormalizedArray(array);
+ }
+ {
+ // map built over sliced keys and items, which land under the entries
struct. Map keys
+ // cannot be null, so this slice avoids the null in `texts`.
+ std::shared_ptr<arrow::Array> offsets =
+ arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 2,
2, 4]").ValueOrDie();
+ std::shared_ptr<arrow::Array> array =
+ arrow::MapArray::FromArrays(offsets, texts->Slice(3, 4),
ints->Slice(4, 4))
+ .ValueOrDie();
+ ASSERT_EQ(array->offset(), 0);
+ const arrow::ArrayData& entries = *array->data()->child_data[0];
+ ASSERT_EQ(entries.offset, 0);
+ ASSERT_EQ(entries.child_data[0]->offset, 3);
+ ASSERT_EQ(entries.child_data[1]->offset, 4);
+ CheckNormalizedArray(array);
+ }
+}
+
+TEST(ArrowUtilsTest, TestNormalizeRecordBatchOffsetsFallsBackForDictionary) {
+ // A dictionary is not part of child_data, so this layout takes the
copying fallback.
+ std::shared_ptr<arrow::Array> indices =
+ arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1,
null, 2, 1, 0]")
+ .ValueOrDie();
+ std::shared_ptr<arrow::Array> dictionary =
+ arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["x",
"yy", "zzz"])")
+ .ValueOrDie();
+ auto dictionary_type = arrow::dictionary(arrow::int32(), arrow::utf8());
+ std::shared_ptr<arrow::Array> array =
+ arrow::DictionaryArray::FromArrays(dictionary_type, indices,
dictionary).ValueOrDie();
+
+ CheckNormalizedSlice(array, /*offset=*/1, /*length=*/4);
+ CheckNormalizedSlice(array, /*offset=*/3, /*length=*/3);
+
+ // The fallback copies, which is also what makes the sharing checks above
meaningful.
+ std::shared_ptr<arrow::RecordBatch> batch = MakeSliceBatch(array,
/*offset=*/1, /*length=*/4);
+ ASSERT_OK_AND_ASSIGN(
+ std::shared_ptr<arrow::RecordBatch> normalized,
+ ArrowUtils::NormalizeRecordBatchOffsets(batch,
arrow::default_memory_pool()));
+ ASSERT_FALSE(IsViewInto(*normalized->column_data(0)->buffers[1],
*array->data()->buffers[1]));
+}
+
TEST(ArrowUtilsTest, TestEqualsIgnoreNullable) {
{
// test simple
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 b3402d4..1b4bdd3 100644
--- a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp
+++ b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp
@@ -397,6 +397,7 @@ Result<std::unique_ptr<arrow::RecordBatchReader>>
PageFilteredRowGroupReader::Re
std::vector<std::shared_ptr<arrow::ChunkedArray>> result_arrays;
result_arrays.reserve(field_indices.size());
+ // TODO(zhouhongfeng.zhf): This loop could be parallelized.
for (int field_idx : field_indices) {
PAIMON_ASSIGN_OR_RAISE(
std::shared_ptr<arrow::ChunkedArray> chunked_array,
@@ -420,7 +421,11 @@ Result<std::unique_ptr<arrow::RecordBatchReader>>
PageFilteredRowGroupReader::Re
field->nullable(),
field->metadata()));
}
auto result_schema = arrow::schema(result_fields);
-
+ // TODO(zhouhongfeng.zhf): This decodes the whole filtered row group up
front, while the
+ // fully-matched path decodes one batch at a time. As a result peak memory
holds every
+ // projected column of the row group instead of a single batch.
+ // Decoding batch by batch would make every returned column single-chunk
so that offset
+ // normalization becomes a no-op.
auto table = arrow::Table::Make(result_schema, std::move(result_arrays),
expected_rows);
return std::make_unique<TableRecordBatchReader>(std::move(table),
max_chunksize, pool);
}