Skip to content

Commit

Permalink
Enable date and time types support in vineyard-graph (#1721)
Browse files Browse the repository at this point in the history 
Fixes #236

Signed-off-by: Tao He <[email protected]>
  • Loading branch information
@sighingnow
sighingnow authored Jan 20, 2024
1 parent b0c3129 commit 3ac1f05
Show file tree
Hide file tree
Showing 18 changed files with 717 additions and 129 deletions.
80 changes: 65 additions & 15 deletions modules/basic/ds/arrow.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ limitations under the License.
#include "client/client.h"
#include "client/ds/blob.h"
#include "common/util/logging.h" // IWYU pragma: keep
#include "common/util/macros.h"

namespace vineyard {

Expand Down Expand Up @@ -108,6 +109,31 @@ class ArrowArrayBuilderVisitor {
return Status::OK();
}

Status Visit(const arrow::Date32Type*) {
builder_ = std::make_shared<Date32Builder>(client_, array_);
return Status::OK();
}

Status Visit(const arrow::Date64Type*) {
builder_ = std::make_shared<Date64Builder>(client_, array_);
return Status::OK();
}

Status Visit(const arrow::Time32Type*) {
builder_ = std::make_shared<Time32Builder>(client_, array_);
return Status::OK();
}

Status Visit(const arrow::Time64Type*) {
builder_ = std::make_shared<Time64Builder>(client_, array_);
return Status::OK();
}

Status Visit(const arrow::TimestampType*) {
builder_ = std::make_shared<TimestampBuilder>(client_, array_);
return Status::OK();
}

Status Visit(const arrow::StringType*) {
builder_ = std::make_shared<StringArrayBuilder>(client_, array_);
return Status::OK();
Expand Down Expand Up @@ -295,10 +321,15 @@ std::shared_ptr<arrow::Array> CastToArray(std::shared_ptr<Object> object) {

template <typename T>
void NumericArray<T>::PostConstruct(const ObjectMeta& meta) {
std::shared_ptr<arrow::DataType> data_type;
if (this->data_type_.empty()) {
data_type = ConvertToArrowType<T>::TypeValue();
} else {
data_type = type_name_to_arrow_type(this->data_type_);
}
this->array_ = std::make_shared<ArrayType>(
ConvertToArrowType<T>::TypeValue(), this->length_,
this->buffer_->ArrowBufferOrEmpty(), this->null_bitmap_->ArrowBuffer(),
this->null_count_, this->offset_);
data_type, this->length_, this->buffer_->ArrowBufferOrEmpty(),
this->null_bitmap_->ArrowBuffer(), this->null_count_, this->offset_);
}

template class NumericArray<int8_t>;
Expand All @@ -311,14 +342,15 @@ template class NumericArray<uint32_t>;
template class NumericArray<uint64_t>;
template class NumericArray<float>;
template class NumericArray<double>;
template class NumericArray<arrow::Date32Type>;
template class NumericArray<arrow::Date64Type>;
template class NumericArray<arrow::Time32Type>;
template class NumericArray<arrow::Time64Type>;
template class NumericArray<arrow::TimestampType>;

template <typename T>
NumericArrayBuilder<T>::NumericArrayBuilder(Client& client)
: NumericArrayBaseBuilder<T>(client) {
std::shared_ptr<ArrayType> array;
CHECK_ARROW_ERROR(ArrowBuilderType<T>{}.Finish(&array));
this->arrays_.emplace_back(array);
}
: NumericArrayBaseBuilder<T>(client) {}

template <typename T>
NumericArrayBuilder<T>::NumericArrayBuilder(
Expand Down Expand Up @@ -353,12 +385,19 @@ template <typename T>
Status NumericArrayBuilder<T>::Build(Client& client) {
memory::VineyardMemoryPool pool(client);
std::shared_ptr<arrow::Array> array;
RETURN_ON_ARROW_ERROR_AND_ASSIGN(
array, arrow_shim::Concatenate(std::move(this->arrays_), &pool));
if (this->arrays_.empty()) {
CHECK_ARROW_ERROR(ArrowBuilderType<T>(ConvertToArrowType<T>::TypeValue(),
arrow::default_memory_pool())
.Finish(&array));
} else {
RETURN_ON_ARROW_ERROR_AND_ASSIGN(
array, arrow_shim::Concatenate(std::move(this->arrays_), &pool));
}
std::shared_ptr<ArrayType> array_ =
std::dynamic_pointer_cast<ArrayType>(array);

this->set_length_(array_->length());
this->set_data_type_(type_name_from_arrow_type(array_->type()));
this->set_null_count_(array_->null_count());
this->set_offset_(array_->offset());
TAKE_BUFFER_OR_NULL_AND_APPLY(this, client, set_buffer_, pool,
Expand All @@ -377,14 +416,19 @@ template class NumericArrayBuilder<uint32_t>;
template class NumericArrayBuilder<uint64_t>;
template class NumericArrayBuilder<float>;
template class NumericArrayBuilder<double>;
template class NumericArrayBuilder<arrow::Date32Type>;
template class NumericArrayBuilder<arrow::Date64Type>;
template class NumericArrayBuilder<arrow::Time32Type>;
template class NumericArrayBuilder<arrow::Time64Type>;
template class NumericArrayBuilder<arrow::TimestampType>;

template <typename T>
FixedNumericArrayBuilder<T>::FixedNumericArrayBuilder(Client& client,
const size_t size)
: NumericArrayBaseBuilder<T>(client), client_(client), size_(size) {
if (size_ > 0) {
VINEYARD_CHECK_OK(client.CreateBlob(size_ * sizeof(T), writer_));
data_ = reinterpret_cast<T*>(writer_->data());
data_ = reinterpret_cast<ArrowValueType<T>*>(writer_->data());
}
}

Expand All @@ -408,7 +452,7 @@ Status FixedNumericArrayBuilder<T>::Make(
out->size_ = size;
if (out->size_ > 0) {
RETURN_ON_ERROR(client.CreateBlob(out->size_ * sizeof(T), out->writer_));
out->data_ = reinterpret_cast<T*>(out->writer_->data());
out->data_ = reinterpret_cast<ArrowValueType<T>*>(out->writer_->data());
}
return Status::OK();
}
Expand All @@ -426,7 +470,7 @@ Status FixedNumericArrayBuilder<T>::Make(
"cannot make builder of size > 0 with a null buffer");
}
out->writer_ = std::move(writer);
out->data_ = reinterpret_cast<T*>(out->writer_->data());
out->data_ = reinterpret_cast<ArrowValueType<T>*>(out->writer_->data());
}
return Status::OK();
}
Expand Down Expand Up @@ -462,15 +506,16 @@ size_t FixedNumericArrayBuilder<T>::size() const {
}

template <typename T>
T* FixedNumericArrayBuilder<T>::MutablePointer(int64_t i) const {
ArrowValueType<T>* FixedNumericArrayBuilder<T>::MutablePointer(
int64_t i) const {
if (data_) {
return data_ + i;
}
return nullptr;
}

template <typename T>
T* FixedNumericArrayBuilder<T>::data() const {
ArrowValueType<T>* FixedNumericArrayBuilder<T>::data() const {
return data_;
}

Expand Down Expand Up @@ -498,6 +543,11 @@ template class FixedNumericArrayBuilder<uint32_t>;
template class FixedNumericArrayBuilder<uint64_t>;
template class FixedNumericArrayBuilder<float>;
template class FixedNumericArrayBuilder<double>;
template class FixedNumericArrayBuilder<arrow::Date32Type>;
template class FixedNumericArrayBuilder<arrow::Date64Type>;
template class FixedNumericArrayBuilder<arrow::Time32Type>;
template class FixedNumericArrayBuilder<arrow::Time64Type>;
template class FixedNumericArrayBuilder<arrow::TimestampType>;

void BooleanArray::PostConstruct(const ObjectMeta& meta) {
this->array_ = std::make_shared<ArrayType>(
Expand Down
16 changes: 13 additions & 3 deletions modules/basic/ds/arrow.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,9 +104,9 @@ class FixedNumericArrayBuilder : public NumericArrayBaseBuilder<T> {

size_t size() const;

T* MutablePointer(int64_t i) const;
ArrowValueType<T>* MutablePointer(int64_t i) const;

T* data() const;
ArrowValueType<T>* data() const;

Status Build(Client& client) override;

Expand All @@ -116,7 +116,7 @@ class FixedNumericArrayBuilder : public NumericArrayBaseBuilder<T> {
Client& client_;
size_t size_ = 0;
std::unique_ptr<BlobWriter> writer_ = nullptr;
T* data_ = nullptr;
ArrowValueType<T>* data_ = nullptr;
};

using Int8Builder = NumericArrayBuilder<int8_t>;
Expand All @@ -129,6 +129,11 @@ using UInt32Builder = NumericArrayBuilder<uint32_t>;
using UInt64Builder = NumericArrayBuilder<uint64_t>;
using FloatBuilder = NumericArrayBuilder<float>;
using DoubleBuilder = NumericArrayBuilder<double>;
using Date32Builder = NumericArrayBuilder<arrow::Date32Type>;
using Date64Builder = NumericArrayBuilder<arrow::Date64Type>;
using Time32Builder = NumericArrayBuilder<arrow::Time32Type>;
using Time64Builder = NumericArrayBuilder<arrow::Time64Type>;
using TimestampBuilder = NumericArrayBuilder<arrow::TimestampType>;

using FixedInt8Builder = FixedNumericArrayBuilder<int8_t>;
using FixedInt16Builder = FixedNumericArrayBuilder<int16_t>;
Expand All @@ -140,6 +145,11 @@ using FixedUInt32Builder = FixedNumericArrayBuilder<uint32_t>;
using FixedUInt64Builder = FixedNumericArrayBuilder<uint64_t>;
using FixedFloatBuilder = FixedNumericArrayBuilder<float>;
using FixedDoubleBuilder = FixedNumericArrayBuilder<double>;
using FixedDate32Builder = FixedNumericArrayBuilder<arrow::Date32Type>;
using FixedDate64Builder = FixedNumericArrayBuilder<arrow::Date64Type>;
using FixedTime32Builder = FixedNumericArrayBuilder<arrow::Time32Type>;
using FixedTime64Builder = FixedNumericArrayBuilder<arrow::Time64Type>;
using FixedTimestampBuilder = FixedNumericArrayBuilder<arrow::TimestampType>;

/**
* @brief BooleanArrayBuilder is designed for constructing Arrow arrays of
Expand Down
8 changes: 7 additions & 1 deletion modules/basic/ds/arrow.vineyard-mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,10 +78,11 @@ class [[vineyard]] NumericArray : public PrimitiveArray,

const size_t length() const { return array_->length(); }

const T* raw_values() const { return array_->raw_values(); }
const ArrowValueType<T>* raw_values() const { return array_->raw_values(); }

private:
[[shared]] size_t length_;
[[shared(optional)]] String data_type_;
[[shared]] int64_t null_count_, offset_;
[[shared]] std::shared_ptr<Blob> buffer_, null_bitmap_;

Expand All @@ -100,6 +101,11 @@ using UInt32Array = NumericArray<uint32_t>;
using UInt64Array = NumericArray<uint64_t>;
using FloatArray = NumericArray<float>;
using DoubleArray = NumericArray<double>;
using Date32Array = NumericArray<arrow::Date32Type>;
using Date64Array = NumericArray<arrow::Date64Type>;
using Time32Array = NumericArray<arrow::Time32Type>;
using Time64Array = NumericArray<arrow::Time64Type>;
using TimestampArray = NumericArray<arrow::TimestampType>;

class BooleanArrayBaseBuilder;

Expand Down
Loading

0 comments on commit 3ac1f05

Please sign in to comment.