Skip to content

Commit cc1e650

Browse files
wgtmacwgtmac
committed
refactor(orc): share validity bitmap append logic
1 parent c3e37ad commit cc1e650

1 file changed

Lines changed: 48 additions & 29 deletions

File tree

‎src/paimon/format/orc/orc_adapter.cpp‎

Lines changed: 48 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -101,10 +101,38 @@ class OrcBackedArrowBuffer : public arrow::ResizableBuffer {
101101
::orc::DataBuffer<T> orc_buffer_;
102102
};
103103

104-
class EmptyBuilder : public arrow::ArrayBuilder {
104+
template <typename BaseBuilder>
105+
class ValidityBitmapAppendingBuilder : public BaseBuilder {
105106
public:
107+
using BaseBuilder::BaseBuilder;
108+
109+
protected:
110+
arrow::Status AppendValidBytesToBitmap(const uint8_t* valid_bytes, int64_t length) {
111+
ARROW_RETURN_NOT_OK(ReserveValidityBitmap(length));
112+
this->UnsafeAppendToBitmap(valid_bytes, length);
113+
return arrow::Status::OK();
114+
}
115+
116+
private:
117+
arrow::Status ReserveValidityBitmap(int64_t additional_capacity) {
118+
int64_t min_capacity = this->length() + additional_capacity;
119+
if (min_capacity <= this->capacity()) {
120+
return arrow::Status::OK();
121+
}
122+
123+
int64_t new_capacity = arrow::BufferBuilder::GrowByFactor(this->capacity(), min_capacity);
124+
ARROW_RETURN_NOT_OK(this->CheckCapacity(new_capacity));
125+
ARROW_RETURN_NOT_OK(this->null_bitmap_builder_.Resize(new_capacity));
126+
this->capacity_ = new_capacity;
127+
return arrow::Status::OK();
128+
}
129+
};
130+
131+
class EmptyBuilder : public ValidityBitmapAppendingBuilder<arrow::ArrayBuilder> {
132+
public:
133+
using Base = ValidityBitmapAppendingBuilder<arrow::ArrayBuilder>;
106134
using arrow::ArrayBuilder::SetNotNull;
107-
explicit EmptyBuilder(arrow::MemoryPool* pool) : arrow::ArrayBuilder(pool) {}
135+
explicit EmptyBuilder(arrow::MemoryPool* pool) : Base(pool) {}
108136
arrow::Status AppendNulls(int64_t length) override {
109137
return arrow::Status::NotImplemented("AppendNulls is not implemented");
110138
}
@@ -144,9 +172,7 @@ class UnPooledBooleanBuilder : public EmptyBuilder {
144172
}
145173

146174
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
147-
ARROW_RETURN_NOT_OK(this->Reserve(length));
148-
this->UnsafeAppendToBitmap(valid_bytes, length);
149-
return arrow::Status::OK();
175+
return this->AppendValidBytesToBitmap(valid_bytes, length);
150176
}
151177

152178
arrow::Status SetData(const uint8_t* data, int64_t length) {
@@ -176,10 +202,12 @@ class UnPooledBooleanBuilder : public EmptyBuilder {
176202
};
177203

178204
template <typename Type>
179-
class UnPooledPrimitiveBuilder : public arrow::NumericBuilder<Type> {
205+
class UnPooledPrimitiveBuilder
206+
: public ValidityBitmapAppendingBuilder<arrow::NumericBuilder<Type>> {
180207
public:
208+
using Base = ValidityBitmapAppendingBuilder<arrow::NumericBuilder<Type>>;
181209
UnPooledPrimitiveBuilder(const std::shared_ptr<arrow::DataType>& type, arrow::MemoryPool* pool)
182-
: arrow::NumericBuilder<Type>(type, pool) {}
210+
: Base(type, pool) {}
183211

184212
void SetData(const std::shared_ptr<arrow::Buffer>& data) {
185213
data_ = data;
@@ -191,9 +219,7 @@ class UnPooledPrimitiveBuilder : public arrow::NumericBuilder<Type> {
191219
}
192220

193221
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
194-
ARROW_RETURN_NOT_OK(this->Reserve(length));
195-
this->UnsafeAppendToBitmap(valid_bytes, length);
196-
return arrow::Status::OK();
222+
return this->AppendValidBytesToBitmap(valid_bytes, length);
197223
}
198224

199225
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
@@ -216,12 +242,14 @@ class UnPooledPrimitiveBuilder : public arrow::NumericBuilder<Type> {
216242
std::shared_ptr<arrow::Buffer> data_;
217243
};
218244

219-
class UnPooledLargeBinaryBuilder : public arrow::LargeBinaryBuilder {
245+
class UnPooledLargeBinaryBuilder
246+
: public ValidityBitmapAppendingBuilder<arrow::LargeBinaryBuilder> {
220247
public:
248+
using Base = ValidityBitmapAppendingBuilder<arrow::LargeBinaryBuilder>;
221249
using arrow::ArrayBuilder::SetNotNull;
222250
UnPooledLargeBinaryBuilder(const std::shared_ptr<arrow::DataType>& type,
223251
arrow::MemoryPool* pool)
224-
: arrow::LargeBinaryBuilder(pool), type_(type) {}
252+
: Base(pool), type_(type) {}
225253

226254
void SetOffsets(const std::shared_ptr<arrow::Buffer>& offsets) {
227255
offsets_ = offsets;
@@ -234,9 +262,7 @@ class UnPooledLargeBinaryBuilder : public arrow::LargeBinaryBuilder {
234262
length_ += length;
235263
}
236264
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
237-
ARROW_RETURN_NOT_OK(this->Reserve(length));
238-
this->UnsafeAppendToBitmap(valid_bytes, length);
239-
return arrow::Status::OK();
265+
return this->AppendValidBytesToBitmap(valid_bytes, length);
240266
}
241267
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
242268
std::shared_ptr<arrow::Buffer> null_bitmap;
@@ -261,11 +287,12 @@ class UnPooledLargeBinaryBuilder : public arrow::LargeBinaryBuilder {
261287
std::shared_ptr<arrow::Buffer> data_;
262288
};
263289

264-
class UnPooledBinaryBuilder : public arrow::BinaryBuilder {
290+
class UnPooledBinaryBuilder : public ValidityBitmapAppendingBuilder<arrow::BinaryBuilder> {
265291
public:
292+
using Base = ValidityBitmapAppendingBuilder<arrow::BinaryBuilder>;
266293
using arrow::ArrayBuilder::SetNotNull;
267294
UnPooledBinaryBuilder(const std::shared_ptr<arrow::DataType>& type, arrow::MemoryPool* pool)
268-
: arrow::BinaryBuilder(pool), type_(type) {}
295+
: Base(pool), type_(type) {}
269296

270297
void SetOffsets(const std::shared_ptr<arrow::Buffer>& offsets) {
271298
offsets_ = offsets;
@@ -281,9 +308,7 @@ class UnPooledBinaryBuilder : public arrow::BinaryBuilder {
281308
}
282309

283310
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
284-
ARROW_RETURN_NOT_OK(this->Reserve(length));
285-
this->UnsafeAppendToBitmap(valid_bytes, length);
286-
return arrow::Status::OK();
311+
return this->AppendValidBytesToBitmap(valid_bytes, length);
287312
}
288313

289314
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
@@ -323,9 +348,7 @@ class UnPooledListBuilder : public EmptyBuilder {
323348
}
324349

325350
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
326-
ARROW_RETURN_NOT_OK(this->Reserve(length));
327-
this->UnsafeAppendToBitmap(valid_bytes, length);
328-
return arrow::Status::OK();
351+
return this->AppendValidBytesToBitmap(valid_bytes, length);
329352
}
330353

331354
void SetOffsets(const std::shared_ptr<arrow::Buffer>& offsets) {
@@ -373,9 +396,7 @@ class UnPooledStructBuilder : public EmptyBuilder {
373396
}
374397

375398
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
376-
ARROW_RETURN_NOT_OK(this->Reserve(length));
377-
this->UnsafeAppendToBitmap(valid_bytes, length);
378-
return arrow::Status::OK();
399+
return this->AppendValidBytesToBitmap(valid_bytes, length);
379400
}
380401

381402
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
@@ -501,9 +522,7 @@ class UnPooledStringDictionaryBuilder : public EmptyBuilder {
501522
indices_ = indices;
502523
}
503524
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
504-
ARROW_RETURN_NOT_OK(this->Reserve(length));
505-
this->UnsafeAppendToBitmap(valid_bytes, length);
506-
return arrow::Status::OK();
525+
return this->AppendValidBytesToBitmap(valid_bytes, length);
507526
}
508527
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
509528
std::shared_ptr<arrow::Buffer> null_bitmap;

0 commit comments

Comments
 (0)