mirror of
https://github.com/wassname/ray.git
synced 2026-07-15 11:25:40 +08:00
Rebase numbuf to latest arrow (#23)
* rebase to latest arrow * fix arrow linking * finish rebase * point arrow to pcmoritz's arrow * fix * fix macOS * fix
This commit is contained in:
committed by
Robert Nishihara
parent
c3db225cbe
commit
c3ab68e88c
@@ -4,21 +4,24 @@ using namespace arrow;
|
||||
|
||||
namespace numbuf {
|
||||
|
||||
std::shared_ptr<arrow::StructArray> DictBuilder::Finish(
|
||||
Status DictBuilder::Finish(
|
||||
std::shared_ptr<Array> key_tuple_data,
|
||||
std::shared_ptr<Array> val_list_data,
|
||||
std::shared_ptr<Array> val_tuple_data,
|
||||
std::shared_ptr<Array> val_dict_data) {
|
||||
std::shared_ptr<Array> val_dict_data,
|
||||
std::shared_ptr<arrow::Array>* out) {
|
||||
// lists and dicts can't be keys of dicts in Python, that is why for
|
||||
// the keys we do not need to collect sublists
|
||||
auto keys = keys_.Finish(nullptr, key_tuple_data, nullptr);
|
||||
auto vals = vals_.Finish(val_list_data, val_tuple_data, val_dict_data);
|
||||
std::shared_ptr<Array> keys, vals;
|
||||
RETURN_NOT_OK(keys_.Finish(nullptr, key_tuple_data, nullptr, &keys));
|
||||
RETURN_NOT_OK(vals_.Finish(val_list_data, val_tuple_data, val_dict_data, &vals));
|
||||
auto keys_field = std::make_shared<Field>("keys", keys->type());
|
||||
auto vals_field = std::make_shared<Field>("vals", vals->type());
|
||||
auto type = std::make_shared<StructType>(std::vector<FieldPtr>({keys_field, vals_field}));
|
||||
std::vector<ArrayPtr> field_arrays({keys, vals});
|
||||
DCHECK(keys->length() == vals->length());
|
||||
return std::make_shared<StructArray>(type, keys->length(), field_arrays);
|
||||
out->reset(new StructArray(type, keys->length(), field_arrays));
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -33,11 +33,12 @@ public:
|
||||
List containing the data from nested dictionaries in the
|
||||
value list of the dictionary
|
||||
*/
|
||||
std::shared_ptr<arrow::StructArray> Finish(
|
||||
arrow::Status Finish(
|
||||
std::shared_ptr<arrow::Array> key_tuple_data,
|
||||
std::shared_ptr<arrow::Array> val_list_data,
|
||||
std::shared_ptr<arrow::Array> val_tuple_data,
|
||||
std::shared_ptr<arrow::Array> val_dict_data);
|
||||
std::shared_ptr<arrow::Array> val_dict_data,
|
||||
std::shared_ptr<arrow::Array>* out);
|
||||
|
||||
private:
|
||||
SequenceBuilder keys_;
|
||||
|
||||
+23
-15
@@ -5,13 +5,16 @@ using namespace arrow;
|
||||
namespace numbuf {
|
||||
|
||||
SequenceBuilder::SequenceBuilder(MemoryPool* pool)
|
||||
: pool_(pool), types_(pool), offsets_(pool),
|
||||
: pool_(pool),
|
||||
types_(pool, std::make_shared<Int8Type>()),
|
||||
offsets_(pool, std::make_shared<Int32Type>()),
|
||||
nones_(pool, std::make_shared<NullType>()),
|
||||
bools_(pool, std::make_shared<BooleanType>()),
|
||||
ints_(pool),
|
||||
ints_(pool, std::make_shared<Int64Type>()),
|
||||
bytes_(pool, std::make_shared<BinaryType>()),
|
||||
strings_(pool, std::make_shared<StringType>()),
|
||||
floats_(pool), doubles_(pool),
|
||||
floats_(pool, std::make_shared<FloatType>()),
|
||||
doubles_(pool, std::make_shared<DoubleType>()),
|
||||
uint8_tensors_(std::make_shared<UInt8Type>(), pool),
|
||||
int8_tensors_(std::make_shared<Int8Type>(), pool),
|
||||
uint16_tensors_(std::make_shared<UInt16Type>(), pool),
|
||||
@@ -76,6 +79,9 @@ Status SequenceBuilder::AppendDouble(double data) {
|
||||
|
||||
#define DEF_TENSOR_APPEND(NAME, TYPE, TAG) \
|
||||
Status SequenceBuilder::AppendTensor(const std::vector<int64_t>& dims, TYPE* data) { \
|
||||
if (TAG == -1) { \
|
||||
NAME.Start(); \
|
||||
} \
|
||||
UPDATE(NAME.length(), TAG); \
|
||||
return NAME.Append(dims, data); \
|
||||
}
|
||||
@@ -109,11 +115,11 @@ Status SequenceBuilder::AppendDict(int32_t size) {
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
#define ADD_ELEMENT(VARNAME, TAG) \
|
||||
if (TAG != -1) { \
|
||||
types[TAG] = VARNAME.type(); \
|
||||
children[TAG] = VARNAME.Finish(); \
|
||||
ARROW_CHECK_OK(nones_.AppendToBitmap(true)); \
|
||||
#define ADD_ELEMENT(VARNAME, TAG) \
|
||||
if (TAG != -1) { \
|
||||
types[TAG] = std::make_shared<Field>("", VARNAME.type()); \
|
||||
RETURN_NOT_OK(VARNAME.Finish(&children[TAG])); \
|
||||
RETURN_NOT_OK(nones_.AppendToBitmap(true)); \
|
||||
}
|
||||
|
||||
#define ADD_SUBSEQUENCE(DATA, OFFSETS, BUILDER, TAG, NAME) \
|
||||
@@ -132,12 +138,13 @@ Status SequenceBuilder::AppendDict(int32_t size) {
|
||||
DCHECK(OFFSETS.size() == 1); \
|
||||
}
|
||||
|
||||
std::shared_ptr<DenseUnionArray> SequenceBuilder::Finish(
|
||||
Status SequenceBuilder::Finish(
|
||||
std::shared_ptr<Array> list_data,
|
||||
std::shared_ptr<Array> tuple_data,
|
||||
std::shared_ptr<Array> dict_data) {
|
||||
std::shared_ptr<Array> dict_data,
|
||||
std::shared_ptr<Array>* out) {
|
||||
|
||||
std::vector<TypePtr> types(num_tags);
|
||||
std::vector<std::shared_ptr<Field>> types(num_tags);
|
||||
std::vector<ArrayPtr> children(num_tags);
|
||||
|
||||
ADD_ELEMENT(bools_, bool_tag);
|
||||
@@ -165,11 +172,12 @@ std::shared_ptr<DenseUnionArray> SequenceBuilder::Finish(
|
||||
ADD_SUBSEQUENCE(tuple_data, tuple_offsets_, tuple_builder, tuple_tag, "tuple");
|
||||
ADD_SUBSEQUENCE(dict_data, dict_offsets_, dict_builder, dict_tag, "dict");
|
||||
|
||||
TypePtr type = TypePtr(new DenseUnionType(types));
|
||||
|
||||
return std::make_shared<DenseUnionArray>(type, types_.length(),
|
||||
std::vector<uint8_t> type_ids = {};
|
||||
TypePtr type = TypePtr(new UnionType(types, type_ids, UnionMode::DENSE));
|
||||
out->reset(new UnionArray(type, types_.length(),
|
||||
children, types_.data(), offsets_.data(),
|
||||
nones_.null_count(), nones_.null_bitmap());
|
||||
nones_.null_count(), nones_.null_bitmap()));
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -79,10 +79,11 @@ class SequenceBuilder {
|
||||
arrow::Status AppendDict(int32_t size);
|
||||
|
||||
//! Finish building the sequence and return the result
|
||||
std::shared_ptr<arrow::DenseUnionArray> Finish(
|
||||
arrow::Status Finish(
|
||||
std::shared_ptr<arrow::Array> list_data,
|
||||
std::shared_ptr<arrow::Array> tuple_data,
|
||||
std::shared_ptr<arrow::Array> dict_data);
|
||||
std::shared_ptr<arrow::Array> dict_data,
|
||||
std::shared_ptr<arrow::Array>* out);
|
||||
|
||||
private:
|
||||
arrow::MemoryPool* pool_;
|
||||
|
||||
@@ -6,19 +6,24 @@ namespace numbuf {
|
||||
|
||||
template<typename T>
|
||||
TensorBuilder<T>::TensorBuilder(const TypePtr& dtype, MemoryPool* pool)
|
||||
: dtype_(dtype) {
|
||||
dim_data_ = std::make_shared<Int64Builder>(pool);
|
||||
dims_ = std::make_shared<ListBuilder>(pool, dim_data_);
|
||||
value_data_ = std::make_shared<PrimitiveBuilder<T>>(pool, dtype);
|
||||
values_ = std::make_shared<ListBuilder>(pool, value_data_);
|
||||
: dtype_(dtype), pool_(pool) {}
|
||||
|
||||
template<typename T>
|
||||
Status TensorBuilder<T>::Start() {
|
||||
dim_data_ = std::make_shared<Int64Builder>(pool_, std::make_shared<Int64Type>());
|
||||
dims_ = std::make_shared<ListBuilder>(pool_, dim_data_);
|
||||
value_data_ = std::make_shared<PrimitiveBuilder<T>>(pool_, dtype_);
|
||||
values_ = std::make_shared<ListBuilder>(pool_, value_data_);
|
||||
auto dims_field = std::make_shared<Field>("dims", dims_->type());
|
||||
auto values_field = std::make_shared<Field>("data", values_->type());
|
||||
auto type = std::make_shared<StructType>(std::vector<FieldPtr>({dims_field, values_field}));
|
||||
tensors_ = std::make_shared<StructBuilder>(pool, type, std::vector<std::shared_ptr<ArrayBuilder>>({dims_, values_}));
|
||||
};
|
||||
tensors_ = std::make_shared<StructBuilder>(pool_, type, std::vector<std::shared_ptr<ArrayBuilder>>({dims_, values_}));
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Status TensorBuilder<T>::Append(const std::vector<int64_t>& dims, const elem_type* data) {
|
||||
DCHECK(tensors_);
|
||||
RETURN_NOT_OK(tensors_->Append());
|
||||
RETURN_NOT_OK(dims_->Append());
|
||||
RETURN_NOT_OK(values_->Append());
|
||||
@@ -32,8 +37,8 @@ Status TensorBuilder<T>::Append(const std::vector<int64_t>& dims, const elem_typ
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
std::shared_ptr<Array> TensorBuilder<T>::Finish() {
|
||||
return tensors_->Finish();
|
||||
Status TensorBuilder<T>::Finish(std::shared_ptr<Array>* out) {
|
||||
return tensors_->Finish(out);
|
||||
}
|
||||
|
||||
template class TensorBuilder<UInt8Type>;
|
||||
|
||||
@@ -18,6 +18,8 @@ public:
|
||||
typedef typename T::c_type elem_type;
|
||||
|
||||
TensorBuilder(const arrow::TypePtr& dtype, arrow::MemoryPool* pool = nullptr);
|
||||
|
||||
arrow::Status Start();
|
||||
|
||||
/*! Append a new tensor.
|
||||
|
||||
@@ -31,7 +33,7 @@ public:
|
||||
arrow::Status Append(const std::vector<int64_t>& dims, const elem_type* data);
|
||||
|
||||
//! Convert the tensors to an Arrow StructArray
|
||||
std::shared_ptr<arrow::Array> Finish();
|
||||
arrow::Status Finish(std::shared_ptr<arrow::Array>* out);
|
||||
|
||||
//! Number of tensors in the column
|
||||
int32_t length() {
|
||||
@@ -44,6 +46,7 @@ public:
|
||||
|
||||
private:
|
||||
arrow::TypePtr dtype_;
|
||||
arrow::MemoryPool* pool_;
|
||||
std::shared_ptr<arrow::Int64Builder> dim_data_;
|
||||
std::shared_ptr<arrow::ListBuilder> dims_;
|
||||
std::shared_ptr<arrow::PrimitiveBuilder<T>> value_data_;
|
||||
|
||||
Reference in New Issue
Block a user