Skip to content

Commit 15ad16c

Browse files
zjw1111zjw1111
authored
refactor: move parquet arrow stream adapters into common utils (alibaba#209)
1 parent f052b2b commit 15ad16c

18 files changed

Lines changed: 203 additions & 137 deletions

‎include/paimon/status.h‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -469,9 +469,4 @@ inline Status GenericToStatus(Status&& st) {
469469
} \
470470
} while (false)
471471

472-
// This is an internal-use macro and should not be used in public headers.
473-
#ifndef RETURN_NOT_OK
474-
#define RETURN_NOT_OK(s) PAIMON_RETURN_NOT_OK(s)
475-
#endif
476-
477472
} // namespace paimon

‎src/paimon/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,8 @@ set(PAIMON_COMMON_SRCS
124124
common/types/data_type_json_parser.cpp
125125
common/types/row_kind.cpp
126126
common/types/row_type.cpp
127+
common/utils/arrow/arrow_input_stream_adapter.cpp
128+
common/utils/arrow/arrow_output_stream_adapter.cpp
127129
common/utils/arrow/arrow_utils.cpp
128130
common/utils/arrow/mem_utils.cpp
129131
common/utils/binary_row_partition_computer.cpp
@@ -432,6 +434,7 @@ if(PAIMON_BUILD_TESTS)
432434
common/types/data_type_test.cpp
433435
common/utils/var_length_int_utils_test.cpp
434436
common/utils/arrow/arrow_utils_test.cpp
437+
common/utils/arrow/arrow_stream_adapter_test.cpp
435438
common/utils/arrow/mem_utils_test.cpp
436439
common/utils/arrow/status_utils_test.cpp
437440
common/utils/concurrent_hash_map_test.cpp

src/paimon/format/parquet/parquet_input_stream_impl.cpp renamed to src/paimon/common/utils/arrow/arrow_input_stream_adapter.cpp

Lines changed: 56 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2024-present Alibaba Inc.
2+
* Copyright 2026-present Alibaba Inc.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -14,80 +14,108 @@
1414
* limitations under the License.
1515
*/
1616

17-
#include "paimon/format/parquet/parquet_input_stream_impl.h"
17+
#include "paimon/common/utils/arrow/arrow_input_stream_adapter.h"
1818

19-
#include <functional>
19+
#include <cstdint>
2020
#include <utility>
2121

2222
#include "arrow/api.h"
23-
#include "arrow/util/future.h"
23+
#include "fmt/format.h"
2424
#include "paimon/common/utils/arrow/status_utils.h"
25+
#include "paimon/common/utils/math.h"
26+
#include "paimon/common/utils/options_utils.h"
2527
#include "paimon/fs/file_system.h"
2628
#include "paimon/macros.h"
2729
#include "paimon/result.h"
2830
#include "paimon/status.h"
2931

30-
namespace arrow::io {
31-
struct IOContext;
32-
} // namespace arrow::io
32+
namespace paimon {
3333

34-
namespace paimon::parquet {
34+
namespace {
3535

36-
ParquetInputStreamImpl::ParquetInputStreamImpl(
37-
const std::shared_ptr<::paimon::InputStream>& input_stream,
36+
template <typename To, typename From>
37+
arrow::Status ValidateArrowIoRange(From value, const char* name) {
38+
if (!InRange<To>(value)) {
39+
return arrow::Status::Invalid(fmt::format("{} value {} is out of bound of type {}", name,
40+
value, OptionsUtils::GetTypeName<To>()));
41+
}
42+
return arrow::Status::OK();
43+
}
44+
45+
} // namespace
46+
47+
ArrowInputStreamAdapter::ArrowInputStreamAdapter(
48+
const std::shared_ptr<paimon::InputStream>& input_stream,
3849
const std::shared_ptr<arrow::MemoryPool>& pool, uint64_t file_size)
3950
: input_stream_(input_stream), pool_(pool), file_size_(file_size) {}
4051

41-
ParquetInputStreamImpl::~ParquetInputStreamImpl() {
52+
ArrowInputStreamAdapter::~ArrowInputStreamAdapter() {
4253
[[maybe_unused]] auto status = DoClose();
4354
}
4455

45-
arrow::Status ParquetInputStreamImpl::Seek(int64_t position) {
56+
arrow::Status ArrowInputStreamAdapter::Seek(int64_t position) {
4657
return ToArrowStatus(input_stream_->Seek(position, SeekOrigin::FS_SEEK_SET));
4758
}
4859

49-
arrow::Result<int64_t> ParquetInputStreamImpl::Read(int64_t nbytes, void* out) {
50-
Result<int32_t> read_bytes = input_stream_->Read(static_cast<char*>(out), nbytes);
60+
arrow::Result<int64_t> ArrowInputStreamAdapter::Read(int64_t nbytes, void* out) {
61+
ARROW_RETURN_NOT_OK(ValidateArrowIoRange<uint32_t>(nbytes, "nbytes"));
62+
Result<int32_t> read_bytes =
63+
input_stream_->Read(static_cast<char*>(out), static_cast<uint32_t>(nbytes));
5164
if (!read_bytes.ok()) {
5265
return ToArrowStatus(read_bytes.status());
5366
}
5467
return read_bytes.value();
5568
}
5669

57-
arrow::Result<std::shared_ptr<arrow::Buffer>> ParquetInputStreamImpl::Read(int64_t nbytes) {
70+
arrow::Result<std::shared_ptr<arrow::Buffer>> ArrowInputStreamAdapter::Read(int64_t nbytes) {
5871
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ResizableBuffer> buffer,
5972
arrow::AllocateResizableBuffer(nbytes, pool_.get()));
6073
ARROW_ASSIGN_OR_RAISE(int64_t read_bytes, Read(nbytes, buffer->mutable_data()));
6174
if (read_bytes < nbytes) {
62-
RETURN_NOT_OK(buffer->Resize(read_bytes));
75+
ARROW_RETURN_NOT_OK(buffer->Resize(read_bytes));
6376
}
6477
return std::shared_ptr<arrow::Buffer>(std::move(buffer));
6578
}
6679

67-
arrow::Result<int64_t> ParquetInputStreamImpl::ReadAt(int64_t position, int64_t nbytes, void* out) {
68-
Result<int32_t> read_bytes = input_stream_->Read(static_cast<char*>(out), nbytes, position);
80+
arrow::Result<int64_t> ArrowInputStreamAdapter::ReadAt(int64_t position, int64_t nbytes,
81+
void* out) {
82+
ARROW_RETURN_NOT_OK(ValidateArrowIoRange<uint64_t>(position, "position"));
83+
ARROW_RETURN_NOT_OK(ValidateArrowIoRange<uint32_t>(nbytes, "nbytes"));
84+
Result<int32_t> read_bytes = input_stream_->Read(
85+
static_cast<char*>(out), static_cast<uint32_t>(nbytes), static_cast<uint64_t>(position));
6986
if (!read_bytes.ok()) {
7087
return ToArrowStatus(read_bytes.status());
7188
}
7289
return read_bytes.value();
7390
}
7491

75-
arrow::Result<std::shared_ptr<arrow::Buffer>> ParquetInputStreamImpl::ReadAt(int64_t position,
76-
int64_t nbytes) {
92+
arrow::Result<std::shared_ptr<arrow::Buffer>> ArrowInputStreamAdapter::ReadAt(int64_t position,
93+
int64_t nbytes) {
7794
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ResizableBuffer> buffer,
7895
arrow::AllocateResizableBuffer(nbytes, pool_.get()));
7996
ARROW_ASSIGN_OR_RAISE(int64_t read_bytes, ReadAt(position, nbytes, buffer->mutable_data()));
8097
if (read_bytes < nbytes) {
81-
RETURN_NOT_OK(buffer->Resize(read_bytes));
98+
ARROW_RETURN_NOT_OK(buffer->Resize(read_bytes));
8299
}
83100
return std::shared_ptr<arrow::Buffer>(std::move(buffer));
84101
}
85102

86-
arrow::Future<std::shared_ptr<arrow::Buffer>> ParquetInputStreamImpl::ReadAsync(
103+
arrow::Future<std::shared_ptr<arrow::Buffer>> ArrowInputStreamAdapter::ReadAsync(
87104
const arrow::io::IOContext& io_context, int64_t position, int64_t nbytes) {
105+
auto fut = arrow::Future<std::shared_ptr<arrow::Buffer>>::Make();
106+
auto range_status = ValidateArrowIoRange<uint64_t>(position, "position");
107+
if (!range_status.ok()) {
108+
fut.MarkFinished(range_status);
109+
return fut;
110+
}
111+
range_status = ValidateArrowIoRange<uint32_t>(nbytes, "nbytes");
112+
if (!range_status.ok()) {
113+
fut.MarkFinished(range_status);
114+
return fut;
115+
}
116+
88117
arrow::Result<std::shared_ptr<arrow::Buffer>> buffer_result =
89118
arrow::AllocateResizableBuffer(nbytes, pool_.get());
90-
auto fut = arrow::Future<std::shared_ptr<arrow::Buffer>>::Make();
91119
if (PAIMON_UNLIKELY(!buffer_result.ok())) {
92120
fut.MarkFinished(buffer_result.status());
93121
return fut;
@@ -105,23 +133,23 @@ arrow::Future<std::shared_ptr<arrow::Buffer>> ParquetInputStreamImpl::ReadAsync(
105133
return fut;
106134
}
107135

108-
arrow::Result<int64_t> ParquetInputStreamImpl::Tell() const {
136+
arrow::Result<int64_t> ArrowInputStreamAdapter::Tell() const {
109137
Result<int64_t> position = input_stream_->GetPos();
110138
if (!position.ok()) {
111139
return ToArrowStatus(position.status());
112140
}
113141
return position.value();
114142
}
115143

116-
arrow::Result<int64_t> ParquetInputStreamImpl::GetSize() {
144+
arrow::Result<int64_t> ArrowInputStreamAdapter::GetSize() {
117145
return static_cast<int64_t>(file_size_);
118146
}
119147

120-
bool ParquetInputStreamImpl::closed() const {
148+
bool ArrowInputStreamAdapter::closed() const {
121149
return closed_;
122150
}
123151

124-
arrow::Status ParquetInputStreamImpl::DoClose() {
152+
arrow::Status ArrowInputStreamAdapter::DoClose() {
125153
if (!closed_) {
126154
Status status = input_stream_->Close();
127155
if (!status.ok()) {
@@ -132,4 +160,4 @@ arrow::Status ParquetInputStreamImpl::DoClose() {
132160
return arrow::Status::OK();
133161
}
134162

135-
} // namespace paimon::parquet
163+
} // namespace paimon

src/paimon/format/parquet/parquet_input_stream_impl.h renamed to src/paimon/common/utils/arrow/arrow_input_stream_adapter.h

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2024-present Alibaba Inc.
2+
* Copyright 2026-present Alibaba Inc.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -22,19 +22,16 @@
2222
#include "arrow/api.h"
2323
#include "arrow/io/interfaces.h"
2424
#include "arrow/util/future.h"
25-
#include "paimon/fs/file_system.h"
25+
#include "paimon/visibility.h"
2626

2727
namespace paimon {
2828
class InputStream;
29-
} // namespace paimon
30-
31-
namespace paimon::parquet {
3229

33-
class ParquetInputStreamImpl : public arrow::io::RandomAccessFile {
30+
class PAIMON_EXPORT ArrowInputStreamAdapter : public arrow::io::RandomAccessFile {
3431
public:
35-
ParquetInputStreamImpl(const std::shared_ptr<::paimon::InputStream>& input_stream,
36-
const std::shared_ptr<arrow::MemoryPool>& pool, uint64_t file_size);
37-
~ParquetInputStreamImpl() override;
32+
ArrowInputStreamAdapter(const std::shared_ptr<paimon::InputStream>& input_stream,
33+
const std::shared_ptr<arrow::MemoryPool>& pool, uint64_t file_size);
34+
~ArrowInputStreamAdapter() override;
3835

3936
// NOTE: In paimon file system definition, position + nbytes should not exceed file_size_.
4037
arrow::Result<int64_t> Read(int64_t nbytes, void* out) override;
@@ -54,10 +51,11 @@ class ParquetInputStreamImpl : public arrow::io::RandomAccessFile {
5451

5552
private:
5653
arrow::Status DoClose();
57-
std::shared_ptr<::paimon::InputStream> input_stream_;
54+
55+
std::shared_ptr<paimon::InputStream> input_stream_;
5856
std::shared_ptr<arrow::MemoryPool> pool_;
5957
uint64_t file_size_;
6058
bool closed_ = false;
6159
};
6260

63-
} // namespace paimon::parquet
61+
} // namespace paimon

src/paimon/format/parquet/parquet_output_stream_impl.cpp renamed to src/paimon/common/utils/arrow/arrow_output_stream_adapter.cpp

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2024-present Alibaba Inc.
2+
* Copyright 2026-present Alibaba Inc.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -14,47 +14,54 @@
1414
* limitations under the License.
1515
*/
1616

17-
#include "paimon/format/parquet/parquet_output_stream_impl.h"
17+
#include "paimon/common/utils/arrow/arrow_output_stream_adapter.h"
1818

1919
#include "arrow/result.h"
20-
#include "arrow/status.h"
20+
#include "fmt/format.h"
2121
#include "paimon/common/utils/arrow/status_utils.h"
22+
#include "paimon/common/utils/math.h"
2223
#include "paimon/fs/file_system.h"
2324
#include "paimon/result.h"
2425

25-
namespace paimon::parquet {
26+
namespace paimon {
2627

27-
ParquetOutputStreamImpl::ParquetOutputStreamImpl(const std::shared_ptr<paimon::OutputStream>& out)
28-
: out_(out), closed_(false) {}
28+
ArrowOutputStreamAdapter::ArrowOutputStreamAdapter(const std::shared_ptr<paimon::OutputStream>& out)
29+
: out_(out) {}
2930

30-
arrow::Status ParquetOutputStreamImpl::Close() {
31-
// output stream close is called by paimon single file writer, no need to close here
31+
arrow::Status ArrowOutputStreamAdapter::Close() {
32+
// output stream close is called by paimon framework(such as single file writer), no need to
33+
// close here
3234
closed_ = true;
3335
return arrow::Status::OK();
3436
}
3537

36-
arrow::Result<int64_t> ParquetOutputStreamImpl::Tell() const {
38+
arrow::Result<int64_t> ArrowOutputStreamAdapter::Tell() const {
3739
paimon::Result<int64_t> pos = out_->GetPos();
3840
if (!pos.ok()) {
3941
return ToArrowStatus(pos.status());
4042
}
4143
return pos.value();
4244
}
4345

44-
bool ParquetOutputStreamImpl::closed() const {
46+
bool ArrowOutputStreamAdapter::closed() const {
4547
return closed_;
4648
}
4749

48-
arrow::Status ParquetOutputStreamImpl::Write(const void* data, int64_t nbytes) {
49-
Result<int32_t> len = out_->Write(static_cast<const char*>(data), nbytes);
50+
arrow::Status ArrowOutputStreamAdapter::Write(const void* data, int64_t nbytes) {
51+
if (!InRange<uint32_t>(nbytes)) {
52+
return arrow::Status::Invalid(
53+
fmt::format("nbytes value {} is out of bound of uint32_t", nbytes));
54+
}
55+
Result<int32_t> len =
56+
out_->Write(static_cast<const char*>(data), static_cast<uint32_t>(nbytes));
5057
if (!len.ok()) {
5158
return ToArrowStatus(len.status());
5259
}
5360
return arrow::Status::OK();
5461
}
5562

56-
arrow::Status ParquetOutputStreamImpl::Flush() {
63+
arrow::Status ArrowOutputStreamAdapter::Flush() {
5764
return ToArrowStatus(out_->Flush());
5865
}
5966

60-
} // namespace paimon::parquet
67+
} // namespace paimon

src/paimon/format/parquet/parquet_output_stream_impl.h renamed to src/paimon/common/utils/arrow/arrow_output_stream_adapter.h

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2024-present Alibaba Inc.
2+
* Copyright 2026-present Alibaba Inc.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -19,21 +19,17 @@
1919
#include <cstdint>
2020
#include <memory>
2121

22-
#include "arrow/buffer.h"
2322
#include "arrow/io/interfaces.h"
2423
#include "arrow/result.h"
2524
#include "arrow/status.h"
26-
#include "paimon/fs/file_system.h"
25+
#include "paimon/visibility.h"
2726

2827
namespace paimon {
2928
class OutputStream;
30-
} // namespace paimon
31-
32-
namespace paimon::parquet {
3329

34-
class ParquetOutputStreamImpl : public arrow::io::OutputStream {
30+
class PAIMON_EXPORT ArrowOutputStreamAdapter : public arrow::io::OutputStream {
3531
public:
36-
explicit ParquetOutputStreamImpl(const std::shared_ptr<paimon::OutputStream>& out);
32+
explicit ArrowOutputStreamAdapter(const std::shared_ptr<paimon::OutputStream>& out);
3733

3834
arrow::Status Close() override;
3935
arrow::Result<int64_t> Tell() const override;
@@ -43,7 +39,7 @@ class ParquetOutputStreamImpl : public arrow::io::OutputStream {
4339

4440
private:
4541
std::shared_ptr<paimon::OutputStream> out_;
46-
bool closed_;
42+
bool closed_ = false;
4743
};
4844

49-
} // namespace paimon::parquet
45+
} // namespace paimon

0 commit comments

Comments
 (0)