Skip to content

Commit e3c38e4

Browse files
wgtmacwgtmac
committed
refactor(orc): use helper for validity bitmap appends
1 parent def2faa commit e3c38e4

1 file changed

Lines changed: 41 additions & 46 deletions

File tree

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

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

104-
template <typename BaseBuilder>
105-
class ValidityBitmapAppendingBuilder : public BaseBuilder {
106-
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();
104+
arrow::Status AppendValidBytesToBitmap(arrow::ArrayBuilder* builder,
105+
arrow::TypedBufferBuilder<bool>* null_bitmap_builder,
106+
int64_t* length, int64_t* null_count,
107+
const uint8_t* valid_bytes, int64_t valid_length) {
108+
int64_t min_capacity = builder->length() + valid_length;
109+
if (min_capacity > builder->capacity()) {
110+
int64_t new_capacity =
111+
arrow::BufferBuilder::GrowByFactor(builder->capacity(), min_capacity);
112+
ARROW_RETURN_NOT_OK(builder->arrow::ArrayBuilder::Resize(new_capacity));
113+
}
114+
115+
if (valid_bytes == nullptr) {
116+
null_bitmap_builder->UnsafeAppend(valid_length, true);
117+
} else {
118+
null_bitmap_builder->UnsafeAppend(valid_bytes, valid_length);
119+
*null_count = null_bitmap_builder->false_count();
128120
}
129-
};
121+
*length += valid_length;
122+
return arrow::Status::OK();
123+
}
130124

131-
class EmptyBuilder : public ValidityBitmapAppendingBuilder<arrow::ArrayBuilder> {
125+
class EmptyBuilder : public arrow::ArrayBuilder {
132126
public:
133-
using Base = ValidityBitmapAppendingBuilder<arrow::ArrayBuilder>;
134127
using arrow::ArrayBuilder::SetNotNull;
135-
explicit EmptyBuilder(arrow::MemoryPool* pool) : Base(pool) {}
128+
explicit EmptyBuilder(arrow::MemoryPool* pool) : arrow::ArrayBuilder(pool) {}
136129
arrow::Status AppendNulls(int64_t length) override {
137130
return arrow::Status::NotImplemented("AppendNulls is not implemented");
138131
}
@@ -172,7 +165,8 @@ class UnPooledBooleanBuilder : public EmptyBuilder {
172165
}
173166

174167
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
175-
return this->AppendValidBytesToBitmap(valid_bytes, length);
168+
return AppendValidBytesToBitmap(this, &null_bitmap_builder_, &length_, &null_count_,
169+
valid_bytes, length);
176170
}
177171

178172
arrow::Status SetData(const uint8_t* data, int64_t length) {
@@ -202,12 +196,10 @@ class UnPooledBooleanBuilder : public EmptyBuilder {
202196
};
203197

204198
template <typename Type>
205-
class UnPooledPrimitiveBuilder
206-
: public ValidityBitmapAppendingBuilder<arrow::NumericBuilder<Type>> {
199+
class UnPooledPrimitiveBuilder : public arrow::NumericBuilder<Type> {
207200
public:
208-
using Base = ValidityBitmapAppendingBuilder<arrow::NumericBuilder<Type>>;
209201
UnPooledPrimitiveBuilder(const std::shared_ptr<arrow::DataType>& type, arrow::MemoryPool* pool)
210-
: Base(type, pool) {}
202+
: arrow::NumericBuilder<Type>(type, pool) {}
211203

212204
void SetData(const std::shared_ptr<arrow::Buffer>& data) {
213205
data_ = data;
@@ -219,7 +211,8 @@ class UnPooledPrimitiveBuilder
219211
}
220212

221213
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
222-
return this->AppendValidBytesToBitmap(valid_bytes, length);
214+
return AppendValidBytesToBitmap(this, &this->null_bitmap_builder_, &this->length_,
215+
&this->null_count_, valid_bytes, length);
223216
}
224217

225218
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
@@ -242,14 +235,12 @@ class UnPooledPrimitiveBuilder
242235
std::shared_ptr<arrow::Buffer> data_;
243236
};
244237

245-
class UnPooledLargeBinaryBuilder
246-
: public ValidityBitmapAppendingBuilder<arrow::LargeBinaryBuilder> {
238+
class UnPooledLargeBinaryBuilder : public arrow::LargeBinaryBuilder {
247239
public:
248-
using Base = ValidityBitmapAppendingBuilder<arrow::LargeBinaryBuilder>;
249240
using arrow::ArrayBuilder::SetNotNull;
250241
UnPooledLargeBinaryBuilder(const std::shared_ptr<arrow::DataType>& type,
251242
arrow::MemoryPool* pool)
252-
: Base(pool), type_(type) {}
243+
: arrow::LargeBinaryBuilder(pool), type_(type) {}
253244

254245
void SetOffsets(const std::shared_ptr<arrow::Buffer>& offsets) {
255246
offsets_ = offsets;
@@ -262,7 +253,8 @@ class UnPooledLargeBinaryBuilder
262253
length_ += length;
263254
}
264255
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
265-
return this->AppendValidBytesToBitmap(valid_bytes, length);
256+
return AppendValidBytesToBitmap(this, &null_bitmap_builder_, &length_, &null_count_,
257+
valid_bytes, length);
266258
}
267259
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
268260
std::shared_ptr<arrow::Buffer> null_bitmap;
@@ -287,12 +279,11 @@ class UnPooledLargeBinaryBuilder
287279
std::shared_ptr<arrow::Buffer> data_;
288280
};
289281

290-
class UnPooledBinaryBuilder : public ValidityBitmapAppendingBuilder<arrow::BinaryBuilder> {
282+
class UnPooledBinaryBuilder : public arrow::BinaryBuilder {
291283
public:
292-
using Base = ValidityBitmapAppendingBuilder<arrow::BinaryBuilder>;
293284
using arrow::ArrayBuilder::SetNotNull;
294285
UnPooledBinaryBuilder(const std::shared_ptr<arrow::DataType>& type, arrow::MemoryPool* pool)
295-
: Base(pool), type_(type) {}
286+
: arrow::BinaryBuilder(pool), type_(type) {}
296287

297288
void SetOffsets(const std::shared_ptr<arrow::Buffer>& offsets) {
298289
offsets_ = offsets;
@@ -308,7 +299,8 @@ class UnPooledBinaryBuilder : public ValidityBitmapAppendingBuilder<arrow::Binar
308299
}
309300

310301
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
311-
return this->AppendValidBytesToBitmap(valid_bytes, length);
302+
return AppendValidBytesToBitmap(this, &null_bitmap_builder_, &length_, &null_count_,
303+
valid_bytes, length);
312304
}
313305

314306
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
@@ -348,7 +340,8 @@ class UnPooledListBuilder : public EmptyBuilder {
348340
}
349341

350342
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
351-
return this->AppendValidBytesToBitmap(valid_bytes, length);
343+
return AppendValidBytesToBitmap(this, &null_bitmap_builder_, &length_, &null_count_,
344+
valid_bytes, length);
352345
}
353346

354347
void SetOffsets(const std::shared_ptr<arrow::Buffer>& offsets) {
@@ -396,7 +389,8 @@ class UnPooledStructBuilder : public EmptyBuilder {
396389
}
397390

398391
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
399-
return this->AppendValidBytesToBitmap(valid_bytes, length);
392+
return AppendValidBytesToBitmap(this, &null_bitmap_builder_, &length_, &null_count_,
393+
valid_bytes, length);
400394
}
401395

402396
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
@@ -522,7 +516,8 @@ class UnPooledStringDictionaryBuilder : public EmptyBuilder {
522516
indices_ = indices;
523517
}
524518
arrow::Status SetNulls(const uint8_t* valid_bytes, int64_t length) {
525-
return this->AppendValidBytesToBitmap(valid_bytes, length);
519+
return AppendValidBytesToBitmap(this, &null_bitmap_builder_, &length_, &null_count_,
520+
valid_bytes, length);
526521
}
527522
arrow::Status FinishInternal(std::shared_ptr<arrow::ArrayData>* out) override {
528523
std::shared_ptr<arrow::Buffer> null_bitmap;

0 commit comments

Comments
 (0)