Merge numbuf.

This commit is contained in:
Philipp Moritz
2016-11-19 12:30:01 -08:00
31 changed files with 2125 additions and 0 deletions
+27
View File
@@ -0,0 +1,27 @@
#include "dict.h"
using namespace arrow;
namespace numbuf {
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<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
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());
out->reset(new StructArray(type, keys->length(), field_arrays));
return Status::OK();
}
}
+50
View File
@@ -0,0 +1,50 @@
#ifndef NUMBUF_DICT_H
#define NUMBUF_DICT_H
#include <arrow/api.h>
#include "sequence.h"
namespace numbuf {
/*! Constructing dictionaries of key/value pairs. Sequences of
keys and values are built separately using a pair of
SequenceBuilders. The resulting Arrow representation
can be obtained via the Finish method.
*/
class DictBuilder {
public:
DictBuilder(arrow::MemoryPool* pool = nullptr)
: keys_(pool), vals_(pool) {}
//! Builder for the keys of the dictionary
SequenceBuilder& keys() { return keys_; }
//! Builder for the values of the dictionary
SequenceBuilder& vals() { return vals_; }
/*! Construct an Arrow StructArray representing the dictionary.
Contains a field "keys" for the keys and "vals" for the values.
\param list_data
List containing the data from nested lists in the value
list of the dictionary
\param dict_data
List containing the data from nested dictionaries in the
value list of the dictionary
*/
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>* out);
private:
SequenceBuilder keys_;
SequenceBuilder vals_;
};
}
#endif
+183
View File
@@ -0,0 +1,183 @@
#include "sequence.h"
using namespace arrow;
namespace numbuf {
SequenceBuilder::SequenceBuilder(MemoryPool* 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, std::make_shared<Int64Type>()),
bytes_(pool, std::make_shared<BinaryType>()),
strings_(pool, std::make_shared<StringType>()),
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),
int16_tensors_(std::make_shared<Int16Type>(), pool),
uint32_tensors_(std::make_shared<UInt32Type>(), pool),
int32_tensors_(std::make_shared<Int32Type>(), pool),
uint64_tensors_(std::make_shared<UInt64Type>(), pool),
int64_tensors_(std::make_shared<Int64Type>(), pool),
float_tensors_(std::make_shared<FloatType>(), pool),
double_tensors_(std::make_shared<DoubleType>(), pool),
list_offsets_({0}), tuple_offsets_({0}), dict_offsets_({0}) {}
#define UPDATE(OFFSET, TAG) \
if (TAG == -1) { \
TAG = num_tags; \
num_tags += 1; \
} \
RETURN_NOT_OK(offsets_.Append(OFFSET)); \
RETURN_NOT_OK(types_.Append(TAG)); \
RETURN_NOT_OK(nones_.AppendToBitmap(true));
Status SequenceBuilder::AppendNone() {
RETURN_NOT_OK(offsets_.Append(0));
RETURN_NOT_OK(types_.Append(0));
return nones_.AppendToBitmap(false);
}
Status SequenceBuilder::AppendBool(bool data) {
UPDATE(bools_.length(), bool_tag);
return bools_.Append(data);
}
Status SequenceBuilder::AppendInt64(int64_t data) {
UPDATE(ints_.length(), int_tag);
return ints_.Append(data);
}
Status SequenceBuilder::AppendUInt64(uint64_t data) {
UPDATE(ints_.length(), int_tag);
return ints_.Append(data);
}
Status SequenceBuilder::AppendBytes(const uint8_t* data, int32_t length) {
UPDATE(bytes_.length(), bytes_tag);
return bytes_.Append(data, length);
}
Status SequenceBuilder::AppendString(const char* data, int32_t length) {
UPDATE(strings_.length(), string_tag);
return strings_.Append(data, length);
}
Status SequenceBuilder::AppendFloat(float data) {
UPDATE(floats_.length(), float_tag);
return floats_.Append(data);
}
Status SequenceBuilder::AppendDouble(double data) {
UPDATE(doubles_.length(), double_tag);
return doubles_.Append(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); \
}
DEF_TENSOR_APPEND(uint8_tensors_, uint8_t, uint8_tensor_tag);
DEF_TENSOR_APPEND(int8_tensors_, int8_t, int8_tensor_tag);
DEF_TENSOR_APPEND(uint16_tensors_, uint16_t, uint16_tensor_tag);
DEF_TENSOR_APPEND(int16_tensors_, int16_t, int16_tensor_tag);
DEF_TENSOR_APPEND(uint32_tensors_, uint32_t, uint32_tensor_tag);
DEF_TENSOR_APPEND(int32_tensors_, int32_t, int32_tensor_tag);
DEF_TENSOR_APPEND(uint64_tensors_, uint64_t, uint64_tensor_tag);
DEF_TENSOR_APPEND(int64_tensors_, int64_t, int64_tensor_tag);
DEF_TENSOR_APPEND(float_tensors_, float, float_tensor_tag);
DEF_TENSOR_APPEND(double_tensors_, double, double_tensor_tag);
Status SequenceBuilder::AppendList(int32_t size) {
UPDATE(list_offsets_.size() - 1, list_tag);
list_offsets_.push_back(list_offsets_.back() + size);
return Status::OK();
}
Status SequenceBuilder::AppendTuple(int32_t size) {
UPDATE(tuple_offsets_.size() - 1, tuple_tag);
tuple_offsets_.push_back(tuple_offsets_.back() + size);
return Status::OK();
}
Status SequenceBuilder::AppendDict(int32_t size) {
UPDATE(dict_offsets_.size() - 1, dict_tag);
dict_offsets_.push_back(dict_offsets_.back() + size);
return Status::OK();
}
#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) \
if (DATA) { \
DCHECK(DATA->length() == OFFSETS.back()); \
auto list_builder = std::make_shared<ListBuilder>(pool_, DATA); \
auto field = std::make_shared<Field>(NAME, list_builder->type()); \
auto type = std::make_shared<StructType>(std::vector<FieldPtr>({field})); \
auto lists = std::vector<std::shared_ptr<ArrayBuilder>>({list_builder}); \
StructBuilder builder(pool_, type, lists); \
OFFSETS.pop_back(); \
ARROW_CHECK_OK(list_builder->Append(OFFSETS.data(), OFFSETS.size())); \
builder.Append(); \
ADD_ELEMENT(builder, TAG); \
} else { \
DCHECK(OFFSETS.size() == 1); \
}
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>* out) {
std::vector<std::shared_ptr<Field>> types(num_tags);
std::vector<ArrayPtr> children(num_tags);
ADD_ELEMENT(bools_, bool_tag);
ADD_ELEMENT(ints_, int_tag);
ADD_ELEMENT(strings_, string_tag);
ADD_ELEMENT(bytes_, bytes_tag);
ADD_ELEMENT(floats_, float_tag);
ADD_ELEMENT(doubles_, double_tag);
ADD_ELEMENT(uint8_tensors_, uint8_tensor_tag);
ADD_ELEMENT(int8_tensors_, int8_tensor_tag);
ADD_ELEMENT(uint16_tensors_, uint16_tensor_tag);
ADD_ELEMENT(int16_tensors_, int16_tensor_tag);
ADD_ELEMENT(uint32_tensors_, uint32_tensor_tag);
ADD_ELEMENT(int32_tensors_, int32_tensor_tag);
ADD_ELEMENT(uint64_tensors_, uint64_tensor_tag);
ADD_ELEMENT(int64_tensors_, int64_tensor_tag);
ADD_ELEMENT(float_tensors_, float_tensor_tag);
ADD_ELEMENT(double_tensors_, double_tensor_tag);
ADD_SUBSEQUENCE(list_data, list_offsets_, list_builder, list_tag, "list");
ADD_SUBSEQUENCE(tuple_data, tuple_offsets_, tuple_builder, tuple_tag, "tuple");
ADD_SUBSEQUENCE(dict_data, dict_offsets_, dict_builder, dict_tag, "dict");
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()));
return Status::OK();
}
}
+144
View File
@@ -0,0 +1,144 @@
#ifndef NUMBUF_LIST_H
#define NUMBUF_LIST_H
#include <arrow/api.h>
#include <arrow/types/union.h>
#include "tensor.h"
namespace numbuf {
/*! A Sequence is a heterogeneous collections of elements. It can contain
scalar Python types, lists, tuples, dictionaries and tensors.
*/
class SequenceBuilder {
public:
SequenceBuilder(arrow::MemoryPool* pool = nullptr);
//! Appending a none to the sequence
arrow::Status AppendNone();
//! Appending a boolean to the sequence
arrow::Status AppendBool(bool data);
//! Appending an int64_t to the sequence
arrow::Status AppendInt64(int64_t data);
//! Appending an uint64_t to the sequence
arrow::Status AppendUInt64(uint64_t data);
//! Append a list of bytes to the sequence
arrow::Status AppendBytes(const uint8_t* data, int32_t length);
//! Appending a string to the sequence
arrow::Status AppendString(const char* data, int32_t length);
//! Appending a float to the sequence
arrow::Status AppendFloat(float data);
//! Appending a double to the sequence
arrow::Status AppendDouble(double data);
/*! Appending a tensor to the sequence
\param dims
A vector of dimensions
\param data
A pointer to the start of the data block. The length of the data block
will be the product of the dimensions
*/
arrow::Status AppendTensor(const std::vector<int64_t>& dims, uint8_t* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, int8_t* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, uint16_t* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, int16_t* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, uint32_t* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, int32_t* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, uint64_t* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, int64_t* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, float* data);
arrow::Status AppendTensor(const std::vector<int64_t>& dims, double* data);
/*! Add a sublist to the sequenc. The data contained in the sublist will be
specified in the "Finish" method.
To construct l = [[11, 22], 33, [44, 55]] you would for example run
list = ListBuilder();
list.AppendList(2);
list.Append(33);
list.AppendList(2);
list.Finish([11, 22, 44, 55]);
list.Finish();
\param size
The size of the sublist
*/
arrow::Status AppendList(int32_t size);
arrow::Status AppendTuple(int32_t size);
arrow::Status AppendDict(int32_t size);
//! Finish building the sequence and return the result
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>* out);
private:
arrow::MemoryPool* pool_;
arrow::Int8Builder types_;
arrow::Int32Builder offsets_;
arrow::NullArrayBuilder nones_;
arrow::BooleanBuilder bools_;
arrow::Int64Builder ints_;
arrow::BinaryBuilder bytes_;
arrow::StringBuilder strings_;
arrow::FloatBuilder floats_;
arrow::DoubleBuilder doubles_;
UInt8TensorBuilder uint8_tensors_;
Int8TensorBuilder int8_tensors_;
UInt16TensorBuilder uint16_tensors_;
Int16TensorBuilder int16_tensors_;
UInt32TensorBuilder uint32_tensors_;
Int32TensorBuilder int32_tensors_;
UInt64TensorBuilder uint64_tensors_;
Int64TensorBuilder int64_tensors_;
FloatTensorBuilder float_tensors_;
DoubleTensorBuilder double_tensors_;
std::vector<int32_t> list_offsets_;
std::vector<int32_t> tuple_offsets_;
std::vector<int32_t> dict_offsets_;
int8_t bool_tag = -1;
int8_t int_tag = -1;
int8_t string_tag = -1;
int8_t bytes_tag = -1;
int8_t float_tag = -1;
int8_t double_tag = -1;
int8_t uint8_tensor_tag = -1;
int8_t int8_tensor_tag = -1;
int8_t uint16_tensor_tag = -1;
int8_t int16_tensor_tag = -1;
int8_t uint32_tensor_tag = -1;
int8_t int32_tensor_tag = -1;
int8_t uint64_tensor_tag = -1;
int8_t int64_tensor_tag = -1;
int8_t float_tensor_tag = -1;
int8_t double_tensor_tag = -1;
int8_t list_tag = -1;
int8_t tuple_tag = -1;
int8_t dict_tag = -1;
int8_t num_tags = 0;
};
} // namespace numbuf
#endif // NUMBUF_LIST_H
+55
View File
@@ -0,0 +1,55 @@
#include "tensor.h"
using namespace arrow;
namespace numbuf {
template<typename T>
TensorBuilder<T>::TensorBuilder(const TypePtr& dtype, MemoryPool* pool)
: 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_}));
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());
int32_t size = 1;
for (auto dim : dims) {
size *= dim;
RETURN_NOT_OK(dim_data_->Append(dim));
}
RETURN_NOT_OK(value_data_->Append(data, size));
return Status::OK(); // tensors_->Append();
}
template<typename T>
Status TensorBuilder<T>::Finish(std::shared_ptr<Array>* out) {
return tensors_->Finish(out);
}
template class TensorBuilder<UInt8Type>;
template class TensorBuilder<Int8Type>;
template class TensorBuilder<UInt16Type>;
template class TensorBuilder<Int16Type>;
template class TensorBuilder<UInt32Type>;
template class TensorBuilder<Int32Type>;
template class TensorBuilder<UInt64Type>;
template class TensorBuilder<Int64Type>;
template class TensorBuilder<FloatType>;
template class TensorBuilder<DoubleType>;
}
+70
View File
@@ -0,0 +1,70 @@
#ifndef NUMBUF_TENSOR_H
#define NUMBUF_TENSOR_H
#include <memory>
#include <arrow/type.h>
#include <arrow/api.h>
namespace numbuf {
/*! This is a class for building a dataframe where each row corresponds to
a Tensor (= multidimensional array) of numerical data. There are two
columns, "dims" which contains an array of dimensions for each Tensor
and "data" which contains data buffer of the Tensor as a flattened array.
*/
template<typename T>
class TensorBuilder {
public:
typedef typename T::c_type elem_type;
TensorBuilder(const arrow::TypePtr& dtype, arrow::MemoryPool* pool = nullptr);
arrow::Status Start();
/*! Append a new tensor.
\param dims
The dimensions of the Tensor
\param data
Pointer to the beginning of the data buffer of the Tensor. The
total length of the buffer is sizeof(elem_type) * product of dims[i] over i
*/
arrow::Status Append(const std::vector<int64_t>& dims, const elem_type* data);
//! Convert the tensors to an Arrow StructArray
arrow::Status Finish(std::shared_ptr<arrow::Array>* out);
//! Number of tensors in the column
int32_t length() {
return tensors_->length();
}
const arrow::TypePtr& type() {
return tensors_->type();
}
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_;
std::shared_ptr<arrow::ListBuilder> values_;
std::shared_ptr<arrow::StructBuilder> tensors_;
};
typedef TensorBuilder<arrow::UInt8Type> UInt8TensorBuilder;
typedef TensorBuilder<arrow::Int8Type> Int8TensorBuilder;
typedef TensorBuilder<arrow::UInt16Type> UInt16TensorBuilder;
typedef TensorBuilder<arrow::Int16Type> Int16TensorBuilder;
typedef TensorBuilder<arrow::UInt32Type> UInt32TensorBuilder;
typedef TensorBuilder<arrow::Int32Type> Int32TensorBuilder;
typedef TensorBuilder<arrow::UInt64Type> UInt64TensorBuilder;
typedef TensorBuilder<arrow::Int64Type> Int64TensorBuilder;
typedef TensorBuilder<arrow::FloatType> FloatTensorBuilder;
typedef TensorBuilder<arrow::DoubleType> DoubleTensorBuilder;
}
#endif // NUMBUF_TENSOR_H