| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| #pragma once |
|
|
| #include <cstddef> |
| #include <cstdint> |
| #include <memory> |
| #include <string> |
| #include <utility> |
| #include <vector> |
|
|
| #include "arrow/buffer.h" |
| #include "arrow/compare.h" |
| #include "arrow/result.h" |
| #include "arrow/status.h" |
| #include "arrow/tensor.h" |
| #include "arrow/type.h" |
| #include "arrow/util/checked_cast.h" |
| #include "arrow/util/macros.h" |
| #include "arrow/util/visibility.h" |
|
|
| namespace arrow { |
|
|
| class MemoryPool; |
|
|
| namespace internal { |
|
|
| ARROW_EXPORT |
| Status CheckSparseIndexMaximumValue(const std::shared_ptr<DataType>& index_value_type, |
| const std::vector<int64_t>& shape); |
|
|
| } |
|
|
| |
| |
|
|
| struct SparseTensorFormat { |
| |
| enum type { |
| |
| COO, |
| |
| CSR, |
| |
| CSC, |
| |
| CSF |
| }; |
| }; |
|
|
| |
| |
| |
| |
| |
| |
| |
| class ARROW_EXPORT SparseIndex { |
| public: |
| explicit SparseIndex(SparseTensorFormat::type format_id) : format_id_(format_id) {} |
|
|
| virtual ~SparseIndex() = default; |
|
|
| |
| SparseTensorFormat::type format_id() const { return format_id_; } |
|
|
| |
| |
| virtual int64_t non_zero_length() const = 0; |
|
|
| |
| virtual std::string ToString() const = 0; |
|
|
| virtual Status ValidateShape(const std::vector<int64_t>& shape) const; |
|
|
| protected: |
| const SparseTensorFormat::type format_id_; |
| }; |
|
|
| namespace internal { |
| template <typename SparseIndexType> |
| class SparseIndexBase : public SparseIndex { |
| public: |
| SparseIndexBase() : SparseIndex(SparseIndexType::format_id) {} |
| }; |
| } |
|
|
| |
| |
|
|
| |
| |
| |
| |
| class ARROW_EXPORT SparseCOOIndex : public internal::SparseIndexBase<SparseCOOIndex> { |
| public: |
| static constexpr SparseTensorFormat::type format_id = SparseTensorFormat::COO; |
|
|
| |
| static Result<std::shared_ptr<SparseCOOIndex>> Make( |
| const std::shared_ptr<Tensor>& coords, bool is_canonical); |
|
|
| |
| static Result<std::shared_ptr<SparseCOOIndex>> Make( |
| const std::shared_ptr<Tensor>& coords); |
|
|
| |
| static Result<std::shared_ptr<SparseCOOIndex>> Make( |
| const std::shared_ptr<DataType>& indices_type, |
| const std::vector<int64_t>& indices_shape, |
| const std::vector<int64_t>& indices_strides, std::shared_ptr<Buffer> indices_data); |
|
|
| |
| static Result<std::shared_ptr<SparseCOOIndex>> Make( |
| const std::shared_ptr<DataType>& indices_type, |
| const std::vector<int64_t>& indices_shape, |
| const std::vector<int64_t>& indices_strides, std::shared_ptr<Buffer> indices_data, |
| bool is_canonical); |
|
|
| |
| |
| |
| |
| |
| static Result<std::shared_ptr<SparseCOOIndex>> Make( |
| const std::shared_ptr<DataType>& indices_type, const std::vector<int64_t>& shape, |
| int64_t non_zero_length, std::shared_ptr<Buffer> indices_data); |
|
|
| |
| |
| |
| |
| static Result<std::shared_ptr<SparseCOOIndex>> Make( |
| const std::shared_ptr<DataType>& indices_type, const std::vector<int64_t>& shape, |
| int64_t non_zero_length, std::shared_ptr<Buffer> indices_data, bool is_canonical); |
|
|
| |
| explicit SparseCOOIndex(const std::shared_ptr<Tensor>& coords, bool is_canonical); |
|
|
| |
| |
| |
| |
| |
| |
| const std::shared_ptr<Tensor>& indices() const { return coords_; } |
|
|
| |
| |
| int64_t non_zero_length() const override { return coords_->shape()[0]; } |
|
|
| |
| |
| |
| bool is_canonical() const { return is_canonical_; } |
|
|
| |
| std::string ToString() const override; |
|
|
| |
| bool Equals(const SparseCOOIndex& other) const { |
| return indices()->Equals(*other.indices()); |
| } |
|
|
| inline Status ValidateShape(const std::vector<int64_t>& shape) const override { |
| ARROW_RETURN_NOT_OK(SparseIndex::ValidateShape(shape)); |
|
|
| if (static_cast<size_t>(coords_->shape()[1]) == shape.size()) { |
| return Status::OK(); |
| } |
|
|
| return Status::Invalid( |
| "shape length is inconsistent with the coords matrix in COO index"); |
| } |
|
|
| protected: |
| std::shared_ptr<Tensor> coords_; |
| bool is_canonical_; |
| }; |
|
|
| namespace internal { |
|
|
| |
| enum class SparseMatrixCompressedAxis : char { |
| |
| ROW, |
| |
| COLUMN |
| }; |
|
|
| ARROW_EXPORT |
| Status ValidateSparseCSXIndex(const std::shared_ptr<DataType>& indptr_type, |
| const std::shared_ptr<DataType>& indices_type, |
| const std::vector<int64_t>& indptr_shape, |
| const std::vector<int64_t>& indices_shape, |
| char const* type_name); |
|
|
| ARROW_EXPORT |
| void CheckSparseCSXIndexValidity(const std::shared_ptr<DataType>& indptr_type, |
| const std::shared_ptr<DataType>& indices_type, |
| const std::vector<int64_t>& indptr_shape, |
| const std::vector<int64_t>& indices_shape, |
| char const* type_name); |
|
|
| template <typename SparseIndexType, SparseMatrixCompressedAxis COMPRESSED_AXIS> |
| class SparseCSXIndex : public SparseIndexBase<SparseIndexType> { |
| public: |
| static constexpr SparseMatrixCompressedAxis kCompressedAxis = COMPRESSED_AXIS; |
|
|
| |
| static Result<std::shared_ptr<SparseIndexType>> Make( |
| const std::shared_ptr<DataType>& indptr_type, |
| const std::shared_ptr<DataType>& indices_type, |
| const std::vector<int64_t>& indptr_shape, const std::vector<int64_t>& indices_shape, |
| std::shared_ptr<Buffer> indptr_data, std::shared_ptr<Buffer> indices_data) { |
| ARROW_RETURN_NOT_OK(ValidateSparseCSXIndex(indptr_type, indices_type, indptr_shape, |
| indices_shape, |
| SparseIndexType::kTypeName)); |
| return std::make_shared<SparseIndexType>( |
| std::make_shared<Tensor>(indptr_type, indptr_data, indptr_shape), |
| std::make_shared<Tensor>(indices_type, indices_data, indices_shape)); |
| } |
|
|
| |
| static Result<std::shared_ptr<SparseIndexType>> Make( |
| const std::shared_ptr<DataType>& indices_type, |
| const std::vector<int64_t>& indptr_shape, const std::vector<int64_t>& indices_shape, |
| std::shared_ptr<Buffer> indptr_data, std::shared_ptr<Buffer> indices_data) { |
| return Make(indices_type, indices_type, indptr_shape, indices_shape, indptr_data, |
| indices_data); |
| } |
|
|
| |
| |
| static Result<std::shared_ptr<SparseIndexType>> Make( |
| const std::shared_ptr<DataType>& indptr_type, |
| const std::shared_ptr<DataType>& indices_type, const std::vector<int64_t>& shape, |
| int64_t non_zero_length, std::shared_ptr<Buffer> indptr_data, |
| std::shared_ptr<Buffer> indices_data) { |
| std::vector<int64_t> indptr_shape({shape[0] + 1}); |
| std::vector<int64_t> indices_shape({non_zero_length}); |
| return Make(indptr_type, indices_type, indptr_shape, indices_shape, indptr_data, |
| indices_data); |
| } |
|
|
| |
| |
| static Result<std::shared_ptr<SparseIndexType>> Make( |
| const std::shared_ptr<DataType>& indices_type, const std::vector<int64_t>& shape, |
| int64_t non_zero_length, std::shared_ptr<Buffer> indptr_data, |
| std::shared_ptr<Buffer> indices_data) { |
| return Make(indices_type, indices_type, shape, non_zero_length, indptr_data, |
| indices_data); |
| } |
|
|
| |
| explicit SparseCSXIndex(const std::shared_ptr<Tensor>& indptr, |
| const std::shared_ptr<Tensor>& indices) |
| : SparseIndexBase<SparseIndexType>(), indptr_(indptr), indices_(indices) { |
| CheckSparseCSXIndexValidity(indptr_->type(), indices_->type(), indptr_->shape(), |
| indices_->shape(), SparseIndexType::kTypeName); |
| } |
|
|
| |
| const std::shared_ptr<Tensor>& indptr() const { return indptr_; } |
|
|
| |
| const std::shared_ptr<Tensor>& indices() const { return indices_; } |
|
|
| |
| |
| int64_t non_zero_length() const override { return indices_->shape()[0]; } |
|
|
| |
| std::string ToString() const override { |
| return std::string(SparseIndexType::kTypeName); |
| } |
|
|
| |
| bool Equals(const SparseIndexType& other) const { |
| return indptr()->Equals(*other.indptr()) && indices()->Equals(*other.indices()); |
| } |
|
|
| inline Status ValidateShape(const std::vector<int64_t>& shape) const override { |
| ARROW_RETURN_NOT_OK(SparseIndex::ValidateShape(shape)); |
|
|
| if (shape.size() < 2) { |
| return Status::Invalid("shape length is too short"); |
| } |
|
|
| if (shape.size() > 2) { |
| return Status::Invalid("shape length is too long"); |
| } |
|
|
| if (indptr_->shape()[0] == shape[static_cast<int64_t>(kCompressedAxis)] + 1) { |
| return Status::OK(); |
| } |
|
|
| return Status::Invalid("shape length is inconsistent with the ", ToString()); |
| } |
|
|
| protected: |
| std::shared_ptr<Tensor> indptr_; |
| std::shared_ptr<Tensor> indices_; |
| }; |
|
|
| } |
|
|
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| class ARROW_EXPORT SparseCSRIndex |
| : public internal::SparseCSXIndex<SparseCSRIndex, |
| internal::SparseMatrixCompressedAxis::ROW> { |
| public: |
| using BaseClass = |
| internal::SparseCSXIndex<SparseCSRIndex, internal::SparseMatrixCompressedAxis::ROW>; |
|
|
| static constexpr SparseTensorFormat::type format_id = SparseTensorFormat::CSR; |
| static constexpr char const* kTypeName = "SparseCSRIndex"; |
|
|
| using SparseCSXIndex::kCompressedAxis; |
| using SparseCSXIndex::Make; |
| using SparseCSXIndex::SparseCSXIndex; |
| }; |
|
|
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| class ARROW_EXPORT SparseCSCIndex |
| : public internal::SparseCSXIndex<SparseCSCIndex, |
| internal::SparseMatrixCompressedAxis::COLUMN> { |
| public: |
| using BaseClass = |
| internal::SparseCSXIndex<SparseCSCIndex, |
| internal::SparseMatrixCompressedAxis::COLUMN>; |
|
|
| static constexpr SparseTensorFormat::type format_id = SparseTensorFormat::CSC; |
| static constexpr char const* kTypeName = "SparseCSCIndex"; |
|
|
| using SparseCSXIndex::kCompressedAxis; |
| using SparseCSXIndex::Make; |
| using SparseCSXIndex::SparseCSXIndex; |
| }; |
|
|
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| class ARROW_EXPORT SparseCSFIndex : public internal::SparseIndexBase<SparseCSFIndex> { |
| public: |
| static constexpr SparseTensorFormat::type format_id = SparseTensorFormat::CSF; |
| static constexpr char const* kTypeName = "SparseCSFIndex"; |
|
|
| |
| static Result<std::shared_ptr<SparseCSFIndex>> Make( |
| const std::shared_ptr<DataType>& indptr_type, |
| const std::shared_ptr<DataType>& indices_type, |
| const std::vector<int64_t>& indices_shapes, const std::vector<int64_t>& axis_order, |
| const std::vector<std::shared_ptr<Buffer>>& indptr_data, |
| const std::vector<std::shared_ptr<Buffer>>& indices_data); |
|
|
| |
| static Result<std::shared_ptr<SparseCSFIndex>> Make( |
| const std::shared_ptr<DataType>& indices_type, |
| const std::vector<int64_t>& indices_shapes, const std::vector<int64_t>& axis_order, |
| const std::vector<std::shared_ptr<Buffer>>& indptr_data, |
| const std::vector<std::shared_ptr<Buffer>>& indices_data) { |
| return Make(indices_type, indices_type, indices_shapes, axis_order, indptr_data, |
| indices_data); |
| } |
|
|
| |
| explicit SparseCSFIndex(const std::vector<std::shared_ptr<Tensor>>& indptr, |
| const std::vector<std::shared_ptr<Tensor>>& indices, |
| const std::vector<int64_t>& axis_order); |
|
|
| |
| const std::vector<std::shared_ptr<Tensor>>& indptr() const { return indptr_; } |
|
|
| |
| const std::vector<std::shared_ptr<Tensor>>& indices() const { return indices_; } |
|
|
| |
| const std::vector<int64_t>& axis_order() const { return axis_order_; } |
|
|
| |
| |
| int64_t non_zero_length() const override { return indices_.back()->shape()[0]; } |
|
|
| |
| std::string ToString() const override; |
|
|
| |
| bool Equals(const SparseCSFIndex& other) const; |
|
|
| protected: |
| std::vector<std::shared_ptr<Tensor>> indptr_; |
| std::vector<std::shared_ptr<Tensor>> indices_; |
| std::vector<int64_t> axis_order_; |
| }; |
|
|
| |
| |
|
|
| |
| class ARROW_EXPORT SparseTensor { |
| public: |
| virtual ~SparseTensor() = default; |
|
|
| SparseTensorFormat::type format_id() const { return sparse_index_->format_id(); } |
|
|
| |
| std::shared_ptr<DataType> type() const { return type_; } |
|
|
| |
| std::shared_ptr<Buffer> data() const { return data_; } |
|
|
| |
| const uint8_t* raw_data() const { return data_->data(); } |
|
|
| |
| uint8_t* raw_mutable_data() const { return data_->mutable_data(); } |
|
|
| |
| const std::vector<int64_t>& shape() const { return shape_; } |
|
|
| |
| const std::shared_ptr<SparseIndex>& sparse_index() const { return sparse_index_; } |
|
|
| |
| int ndim() const { return static_cast<int>(shape_.size()); } |
|
|
| |
| const std::vector<std::string>& dim_names() const { return dim_names_; } |
|
|
| |
| const std::string& dim_name(int i) const; |
|
|
| |
| int64_t size() const; |
|
|
| |
| bool is_mutable() const { return data_->is_mutable(); } |
|
|
| |
| int64_t non_zero_length() const { |
| return sparse_index_ ? sparse_index_->non_zero_length() : 0; |
| } |
|
|
| |
| bool Equals(const SparseTensor& other, |
| const EqualOptions& = EqualOptions::Defaults()) const; |
|
|
| |
| |
| |
| Result<std::shared_ptr<Tensor>> ToTensor(MemoryPool* pool) const; |
| Result<std::shared_ptr<Tensor>> ToTensor() const { |
| return ToTensor(default_memory_pool()); |
| } |
|
|
| protected: |
| |
| SparseTensor(const std::shared_ptr<DataType>& type, const std::shared_ptr<Buffer>& data, |
| const std::vector<int64_t>& shape, |
| const std::shared_ptr<SparseIndex>& sparse_index, |
| const std::vector<std::string>& dim_names); |
|
|
| std::shared_ptr<DataType> type_; |
| std::shared_ptr<Buffer> data_; |
| std::vector<int64_t> shape_; |
| std::shared_ptr<SparseIndex> sparse_index_; |
|
|
| |
| std::vector<std::string> dim_names_; |
| }; |
|
|
| |
| |
|
|
| namespace internal { |
|
|
| ARROW_EXPORT |
| Status MakeSparseTensorFromTensor(const Tensor& tensor, |
| SparseTensorFormat::type sparse_format_id, |
| const std::shared_ptr<DataType>& index_value_type, |
| MemoryPool* pool, |
| std::shared_ptr<SparseIndex>* out_sparse_index, |
| std::shared_ptr<Buffer>* out_data); |
|
|
| } |
|
|
| |
| |
| template <typename SparseIndexType> |
| class SparseTensorImpl : public SparseTensor { |
| public: |
| virtual ~SparseTensorImpl() = default; |
|
|
| |
| SparseTensorImpl(const std::shared_ptr<SparseIndexType>& sparse_index, |
| const std::shared_ptr<DataType>& type, |
| const std::shared_ptr<Buffer>& data, const std::vector<int64_t>& shape, |
| const std::vector<std::string>& dim_names) |
| : SparseTensor(type, data, shape, sparse_index, dim_names) {} |
|
|
| |
| SparseTensorImpl(const std::shared_ptr<DataType>& type, |
| const std::vector<int64_t>& shape, |
| const std::vector<std::string>& dim_names = {}) |
| : SparseTensorImpl(NULLPTR, type, NULLPTR, shape, dim_names) {} |
|
|
| |
| static inline Result<std::shared_ptr<SparseTensorImpl<SparseIndexType>>> Make( |
| const std::shared_ptr<SparseIndexType>& sparse_index, |
| const std::shared_ptr<DataType>& type, const std::shared_ptr<Buffer>& data, |
| const std::vector<int64_t>& shape, const std::vector<std::string>& dim_names) { |
| if (!is_tensor_supported(type->id())) { |
| return Status::Invalid(type->ToString(), |
| " is not valid data type for a sparse tensor"); |
| } |
| ARROW_RETURN_NOT_OK(sparse_index->ValidateShape(shape)); |
| if (dim_names.size() > 0 && dim_names.size() != shape.size()) { |
| return Status::Invalid("dim_names length is inconsistent with shape"); |
| } |
| return std::make_shared<SparseTensorImpl<SparseIndexType>>(sparse_index, type, data, |
| shape, dim_names); |
| } |
|
|
| |
| |
| |
| |
| static inline Result<std::shared_ptr<SparseTensorImpl<SparseIndexType>>> Make( |
| const Tensor& tensor, const std::shared_ptr<DataType>& index_value_type, |
| MemoryPool* pool = default_memory_pool()) { |
| std::shared_ptr<SparseIndex> sparse_index; |
| std::shared_ptr<Buffer> data; |
| ARROW_RETURN_NOT_OK(internal::MakeSparseTensorFromTensor( |
| tensor, SparseIndexType::format_id, index_value_type, pool, &sparse_index, |
| &data)); |
| return std::make_shared<SparseTensorImpl<SparseIndexType>>( |
| internal::checked_pointer_cast<SparseIndexType>(sparse_index), tensor.type(), |
| data, tensor.shape(), tensor.dim_names_); |
| } |
|
|
| static inline Result<std::shared_ptr<SparseTensorImpl<SparseIndexType>>> Make( |
| const Tensor& tensor, MemoryPool* pool = default_memory_pool()) { |
| return Make(tensor, int64(), pool); |
| } |
|
|
| private: |
| ARROW_DISALLOW_COPY_AND_ASSIGN(SparseTensorImpl); |
| }; |
|
|
| |
| using SparseCOOTensor = SparseTensorImpl<SparseCOOIndex>; |
|
|
| |
| using SparseCSRMatrix = SparseTensorImpl<SparseCSRIndex>; |
|
|
| |
| using SparseCSCMatrix = SparseTensorImpl<SparseCSCIndex>; |
|
|
| |
| using SparseCSFTensor = SparseTensorImpl<SparseCSFIndex>; |
|
|
| } |
|
|