Skip to content

Commit 57c525f

Browse files
rokwesm
rok
authored and
wesm
committed
ARROW-7428: [Format][C++] Add serialization for CSF sparse tensors
This is to solve [ARROW-7428](https://issues.apache.org/jira/browse/ARROW-7428). Closes apache#6340 from rok/ARROW-7428 Authored-by: Rok <rok@mihevc.org> Signed-off-by: Wes McKinney <wesm+git@apache.org>
1 parent 58c0941 commit 57c525f

7 files changed

Lines changed: 541 additions & 7 deletions

File tree

‎cpp/src/arrow/ipc/metadata_internal.cc‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,14 @@
2929
#include "arrow/io/interfaces.h"
3030
#include "arrow/ipc/dictionary.h"
3131
#include "arrow/ipc/message.h"
32+
#include "arrow/ipc/util.h"
3233
#include "arrow/sparse_tensor.h"
3334
#include "arrow/status.h"
3435
#include "arrow/type.h"
3536
#include "arrow/type_traits.h"
3637
#include "arrow/util/checked_cast.h"
3738
#include "arrow/util/key_value_metadata.h"
39+
#include "arrow/util/ubsan.h"
3840
#include "arrow/visitor_inline.h"
3941

4042
#include "generated/File_generated.h"
@@ -932,6 +934,64 @@ Status MakeSparseMatrixIndexCSX(FBB& fbb, const SparseIndexType& sparse_index,
932934
return Status::OK();
933935
}
934936

937+
Status MakeSparseTensorIndexCSF(FBB& fbb, const SparseCSFIndex& sparse_index,
938+
const std::vector<BufferMetadata>& buffers,
939+
flatbuf::SparseTensorIndex* fb_sparse_index_type,
940+
Offset* fb_sparse_index, size_t* num_buffers) {
941+
*fb_sparse_index_type = flatbuf::SparseTensorIndex::SparseTensorIndexCSF;
942+
const int ndim = static_cast<int>(sparse_index.axis_order().size());
943+
944+
// We assume that the value type of indptr tensor is an integer.
945+
const auto& indptr_value_type =
946+
checked_cast<const IntegerType&>(*sparse_index.indptr()[0]->type());
947+
auto indptr_type_offset = flatbuf::CreateInt(fbb, indptr_value_type.bit_width(),
948+
indptr_value_type.is_signed());
949+
950+
// We assume that the value type of indices tensor is an integer.
951+
const auto& indices_value_type =
952+
checked_cast<const IntegerType&>(*sparse_index.indices()[0]->type());
953+
auto indices_type_offset = flatbuf::CreateInt(fbb, indices_value_type.bit_width(),
954+
indices_value_type.is_signed());
955+
956+
const int64_t indptr_elem_size = indptr_value_type.bit_width() / 8;
957+
const int64_t indices_elem_size = indices_value_type.bit_width() / 8;
958+
959+
int64_t offset = 0;
960+
std::vector<flatbuf::Buffer> indptr, indices;
961+
962+
for (const std::shared_ptr<arrow::Tensor> tensor : sparse_index.indptr()) {
963+
const int64_t size = tensor->data()->size() / indptr_elem_size;
964+
const int64_t padded_size = PaddedLength(tensor->data()->size(), kArrowIpcAlignment);
965+
966+
indptr.push_back({offset, size});
967+
offset += padded_size;
968+
}
969+
for (const std::shared_ptr<arrow::Tensor> tensor : sparse_index.indices()) {
970+
const int64_t size = tensor->data()->size() / indices_elem_size;
971+
const int64_t padded_size = PaddedLength(tensor->data()->size(), kArrowIpcAlignment);
972+
973+
indices.push_back({offset, size});
974+
offset += padded_size;
975+
}
976+
977+
auto fb_indices = fbb.CreateVectorOfStructs(indices);
978+
auto fb_indptr = fbb.CreateVectorOfStructs(indptr);
979+
980+
std::vector<int> axis_order;
981+
for (int i = 0; i < ndim; ++i) {
982+
axis_order.emplace_back(static_cast<int>(sparse_index.axis_order()[i]));
983+
}
984+
auto fb_axis_order =
985+
fbb.CreateVector(arrow::util::MakeNonNull(axis_order.data()), axis_order.size());
986+
987+
*fb_sparse_index =
988+
flatbuf::CreateSparseTensorIndexCSF(fbb, indptr_type_offset, fb_indptr,
989+
indices_type_offset, fb_indices, fb_axis_order)
990+
.Union();
991+
*num_buffers = 2 * ndim - 1;
992+
return Status::OK();
993+
}
994+
935995
Status MakeSparseTensorIndex(FBB& fbb, const SparseIndex& sparse_index,
936996
const std::vector<BufferMetadata>& buffers,
937997
flatbuf::SparseTensorIndex* fb_sparse_index_type,
@@ -955,7 +1015,14 @@ Status MakeSparseTensorIndex(FBB& fbb, const SparseIndex& sparse_index,
9551015
fb_sparse_index_type, fb_sparse_index, num_buffers));
9561016
break;
9571017

1018+
case SparseTensorFormat::CSF:
1019+
RETURN_NOT_OK(MakeSparseTensorIndexCSF(
1020+
fbb, checked_cast<const SparseCSFIndex&>(sparse_index), buffers,
1021+
fb_sparse_index_type, fb_sparse_index, num_buffers));
1022+
break;
1023+
9581024
default:
1025+
*fb_sparse_index_type = flatbuf::SparseTensorIndex::NONE; // Silence warnings
9591026
std::stringstream ss;
9601027
ss << "Unsupported sparse tensor format:: " << sparse_index.ToString() << std::endl;
9611028
return Status::NotImplemented(ss.str());
@@ -1215,6 +1282,23 @@ Status GetSparseCSXIndexMetadata(const flatbuf::SparseMatrixIndexCSX* sparse_ind
12151282
return Status::OK();
12161283
}
12171284

1285+
Status GetSparseCSFIndexMetadata(const flatbuf::SparseTensorIndexCSF* sparse_index,
1286+
std::vector<int64_t>* axis_order,
1287+
std::vector<int64_t>* indices_size,
1288+
std::shared_ptr<DataType>* indptr_type,
1289+
std::shared_ptr<DataType>* indices_type) {
1290+
RETURN_NOT_OK(IntFromFlatbuffer(sparse_index->indptrType(), indptr_type));
1291+
RETURN_NOT_OK(IntFromFlatbuffer(sparse_index->indicesType(), indices_type));
1292+
1293+
const int ndim = static_cast<int>(sparse_index->axisOrder()->size());
1294+
for (int i = 0; i < ndim; ++i) {
1295+
axis_order->push_back(sparse_index->axisOrder()->Get(i));
1296+
indices_size->push_back(sparse_index->indicesBuffers()->Get(i)->length());
1297+
}
1298+
1299+
return Status::OK();
1300+
}
1301+
12181302
Status GetSparseTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>* type,
12191303
std::vector<int64_t>* shape,
12201304
std::vector<std::string>* dim_names,
@@ -1269,6 +1353,10 @@ Status GetSparseTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>
12691353
}
12701354
} break;
12711355

1356+
case flatbuf::SparseTensorIndex::SparseTensorIndexCSF:
1357+
*sparse_tensor_format_id = SparseTensorFormat::CSF;
1358+
break;
1359+
12721360
default:
12731361
return Status::Invalid("Unrecognized sparse index type");
12741362
}

‎cpp/src/arrow/ipc/metadata_internal.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,13 @@ Status GetSparseCSXIndexMetadata(const flatbuf::SparseMatrixIndexCSX* sparse_ind
137137
std::shared_ptr<DataType>* indptr_type,
138138
std::shared_ptr<DataType>* indices_type);
139139

140+
// EXPERIMENTAL: Extracting metadata of a SparseCSFIndex from the message
141+
Status GetSparseCSFIndexMetadata(const flatbuf::SparseTensorIndexCSF* sparse_index,
142+
std::vector<int64_t>* axis_order,
143+
std::vector<int64_t>* indices_size,
144+
std::shared_ptr<DataType>* indptr_type,
145+
std::shared_ptr<DataType>* indices_type);
146+
140147
// EXPERIMENTAL: Extracting metadata of a sparse tensor from the message
141148
Status GetSparseTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>* type,
142149
std::vector<int64_t>* shape,

‎cpp/src/arrow/ipc/read_write_test.cc‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1402,6 +1402,66 @@ class TestSparseTensorRoundTrip : public ::testing::Test, public IpcTestFixture
14021402
ASSERT_TRUE(result->Equals(sparse_tensor));
14031403
}
14041404

1405+
void CheckSparseCSFTensorRoundTrip(const SparseCSFTensor& sparse_tensor) {
1406+
const auto& type = checked_cast<const FixedWidthType&>(*sparse_tensor.type());
1407+
const int elem_size = type.bit_width() / 8;
1408+
const int index_elem_size = sizeof(typename IndexValueType::c_type);
1409+
1410+
int32_t metadata_length;
1411+
int64_t body_length;
1412+
1413+
ASSERT_OK(mmap_->Seek(0));
1414+
1415+
ASSERT_OK(
1416+
WriteSparseTensor(sparse_tensor, mmap_.get(), &metadata_length, &body_length));
1417+
1418+
const auto& sparse_index =
1419+
checked_cast<const SparseCSFIndex&>(*sparse_tensor.sparse_index());
1420+
1421+
const int64_t ndim = sparse_index.axis_order().size();
1422+
int64_t indptr_length = 0;
1423+
int64_t indices_length = 0;
1424+
1425+
for (int64_t i = 0; i < ndim - 1; ++i) {
1426+
indptr_length += BitUtil::RoundUpToMultipleOf8(index_elem_size *
1427+
sparse_index.indptr()[i]->size());
1428+
}
1429+
for (int64_t i = 0; i < ndim; ++i) {
1430+
indices_length += BitUtil::RoundUpToMultipleOf8(index_elem_size *
1431+
sparse_index.indices()[i]->size());
1432+
}
1433+
const int64_t data_length =
1434+
BitUtil::RoundUpToMultipleOf8(elem_size * sparse_tensor.non_zero_length());
1435+
const int64_t expected_body_length = indptr_length + indices_length + data_length;
1436+
ASSERT_EQ(expected_body_length, body_length);
1437+
1438+
ASSERT_OK(mmap_->Seek(0));
1439+
1440+
std::shared_ptr<SparseTensor> result;
1441+
ASSERT_OK_AND_ASSIGN(result, ReadSparseTensor(mmap_.get()));
1442+
ASSERT_EQ(SparseTensorFormat::CSF, result->format_id());
1443+
1444+
const auto& resulted_sparse_index =
1445+
checked_cast<const SparseCSFIndex&>(*result->sparse_index());
1446+
1447+
int64_t out_indptr_length = 0;
1448+
int64_t out_indices_length = 0;
1449+
for (int i = 0; i < ndim - 1; ++i) {
1450+
out_indptr_length += BitUtil::RoundUpToMultipleOf8(
1451+
index_elem_size * resulted_sparse_index.indptr()[i]->size());
1452+
}
1453+
for (int i = 0; i < ndim; ++i) {
1454+
out_indices_length += BitUtil::RoundUpToMultipleOf8(
1455+
index_elem_size * resulted_sparse_index.indices()[i]->size());
1456+
}
1457+
1458+
ASSERT_EQ(out_indptr_length, indptr_length);
1459+
ASSERT_EQ(out_indices_length, indices_length);
1460+
ASSERT_EQ(result->data()->size(), data_length);
1461+
ASSERT_TRUE(resulted_sparse_index.Equals(sparse_index));
1462+
ASSERT_TRUE(result->Equals(sparse_tensor));
1463+
}
1464+
14051465
protected:
14061466
std::shared_ptr<SparseCOOIndex> MakeSparseCOOIndex(
14071467
const std::vector<int64_t>& coords_shape,
@@ -1564,9 +1624,30 @@ TYPED_TEST_P(TestSparseTensorRoundTrip, WithSparseCSCIndex) {
15641624
this->CheckSparseCSXMatrixRoundTrip(*st);
15651625
}
15661626

1627+
TYPED_TEST_P(TestSparseTensorRoundTrip, WithSparseCSFIndex) {
1628+
using IndexValueType = TypeParam;
1629+
1630+
std::string path = "test-write-sparse-csf-tensor";
1631+
constexpr int64_t kBufferSize = 1 << 20;
1632+
ASSERT_OK_AND_ASSIGN(this->mmap_,
1633+
io::MemoryMapFixture::InitMemoryMap(kBufferSize, path));
1634+
1635+
std::vector<int64_t> shape = {4, 6};
1636+
std::vector<std::string> dim_names = {"foo", "bar", "baz"};
1637+
std::vector<int64_t> values = {1, 0, 2, 0, 0, 3, 0, 4, 5, 0, 6, 0,
1638+
0, 11, 0, 12, 13, 0, 14, 0, 0, 15, 0, 16};
1639+
1640+
auto data = Buffer::Wrap(values);
1641+
NumericTensor<Int64Type> t(data, shape, {}, dim_names);
1642+
std::shared_ptr<SparseCSFTensor> st;
1643+
ASSERT_OK_AND_ASSIGN(
1644+
st, SparseCSFTensor::Make(t, TypeTraits<IndexValueType>::type_singleton()));
1645+
1646+
this->CheckSparseCSFTensorRoundTrip(*st);
1647+
}
15671648
REGISTER_TYPED_TEST_SUITE_P(TestSparseTensorRoundTrip, WithSparseCOOIndexRowMajor,
15681649
WithSparseCOOIndexColumnMajor, WithSparseCSRIndex,
1569-
WithSparseCSCIndex);
1650+
WithSparseCSCIndex, WithSparseCSFIndex);
15701651

15711652
INSTANTIATE_TYPED_TEST_SUITE_P(TestInt8, TestSparseTensorRoundTrip, Int8Type);
15721653
INSTANTIATE_TYPED_TEST_SUITE_P(TestUInt8, TestSparseTensorRoundTrip, UInt8Type);

‎cpp/src/arrow/ipc/reader.cc‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1058,6 +1058,35 @@ Result<std::shared_ptr<SparseIndex>> ReadSparseCSXIndex(
10581058
}
10591059
}
10601060

1061+
Result<std::shared_ptr<SparseIndex>> ReadSparseCSFIndex(
1062+
const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape,
1063+
io::RandomAccessFile* file) {
1064+
auto* sparse_index = sparse_tensor->sparseIndex_as_SparseTensorIndexCSF();
1065+
const auto ndim = static_cast<int64_t>(shape.size());
1066+
auto* indptr_buffers = sparse_index->indptrBuffers();
1067+
auto* indices_buffers = sparse_index->indicesBuffers();
1068+
std::vector<std::shared_ptr<Buffer>> indptr_data(ndim - 1);
1069+
std::vector<std::shared_ptr<Buffer>> indices_data(ndim);
1070+
1071+
std::shared_ptr<DataType> indptr_type, indices_type;
1072+
std::vector<int64_t> axis_order, indices_size;
1073+
1074+
RETURN_NOT_OK(internal::GetSparseCSFIndexMetadata(
1075+
sparse_index, &axis_order, &indices_size, &indptr_type, &indices_type));
1076+
for (int i = 0; i < static_cast<int>(indptr_buffers->Length()); ++i) {
1077+
ARROW_ASSIGN_OR_RAISE(indptr_data[i], file->ReadAt(indptr_buffers->Get(i)->offset(),
1078+
indptr_buffers->Get(i)->length()));
1079+
}
1080+
for (int i = 0; i < static_cast<int>(indices_buffers->Length()); ++i) {
1081+
ARROW_ASSIGN_OR_RAISE(indices_data[i],
1082+
file->ReadAt(indices_buffers->Get(i)->offset(),
1083+
indices_buffers->Get(i)->length()));
1084+
}
1085+
1086+
return SparseCSFIndex::Make(indptr_type, indices_type, indices_size, axis_order,
1087+
indptr_data, indices_data);
1088+
}
1089+
10611090
Result<std::shared_ptr<SparseTensor>> MakeSparseTensorWithSparseCOOIndex(
10621091
const std::shared_ptr<DataType>& type, const std::vector<int64_t>& shape,
10631092
const std::vector<std::string>& dim_names,
@@ -1082,6 +1111,14 @@ Result<std::shared_ptr<SparseTensor>> MakeSparseTensorWithSparseCSCIndex(
10821111
return SparseCSCMatrix::Make(sparse_index, type, data, shape, dim_names);
10831112
}
10841113

1114+
Result<std::shared_ptr<SparseTensor>> MakeSparseTensorWithSparseCSFIndex(
1115+
const std::shared_ptr<DataType>& type, const std::vector<int64_t>& shape,
1116+
const std::vector<std::string>& dim_names,
1117+
const std::shared_ptr<SparseCSFIndex>& sparse_index,
1118+
const std::shared_ptr<Buffer>& data) {
1119+
return SparseCSFTensor::Make(sparse_index, type, data, shape, dim_names);
1120+
}
1121+
10851122
Status ReadSparseTensorMetadata(const Buffer& metadata,
10861123
std::shared_ptr<DataType>* out_type,
10871124
std::vector<int64_t>* out_shape,
@@ -1131,6 +1168,9 @@ Result<size_t> GetSparseTensorBodyBufferCount(SparseTensorFormat::type format_id
11311168
case SparseTensorFormat::CSC:
11321169
return 3;
11331170

1171+
case SparseTensorFormat::CSF:
1172+
return 3;
1173+
11341174
default:
11351175
return Status::Invalid("Unrecognized sparse tensor format");
11361176
}
@@ -1215,6 +1255,33 @@ Result<std::shared_ptr<SparseTensor>> ReadSparseTensorPayload(const IpcPayload&
12151255
return MakeSparseTensorWithSparseCSCIndex(type, shape, dim_names, sparse_index,
12161256
non_zero_length, payload.body_buffers[2]);
12171257
}
1258+
case SparseTensorFormat::CSF: {
1259+
std::shared_ptr<SparseCSFIndex> sparse_index;
1260+
std::shared_ptr<DataType> indptr_type, indices_type;
1261+
std::vector<int64_t> axis_order, indices_size;
1262+
1263+
RETURN_NOT_OK(internal::GetSparseCSFIndexMetadata(
1264+
sparse_tensor->sparseIndex_as_SparseTensorIndexCSF(), &axis_order,
1265+
&indices_size, &indptr_type, &indices_type));
1266+
ARROW_CHECK_EQ(indptr_type, indices_type);
1267+
1268+
const int64_t ndim = shape.size();
1269+
std::vector<std::shared_ptr<Buffer>> indptr_data(ndim - 1);
1270+
std::vector<std::shared_ptr<Buffer>> indices_data(ndim);
1271+
1272+
for (int64_t i = 0; i < ndim - 1; ++i) {
1273+
indptr_data[i] = payload.body_buffers[i];
1274+
}
1275+
for (int64_t i = 0; i < ndim; ++i) {
1276+
indices_data[i] = payload.body_buffers[i + ndim - 1];
1277+
}
1278+
1279+
ARROW_ASSIGN_OR_RAISE(sparse_index,
1280+
SparseCSFIndex::Make(indptr_type, indices_type, indices_size,
1281+
axis_order, indptr_data, indices_data));
1282+
return MakeSparseTensorWithSparseCSFIndex(type, shape, dim_names, sparse_index,
1283+
payload.body_buffers[2 * ndim - 1]);
1284+
}
12181285
default:
12191286
return Status::Invalid("Unsupported sparse index format");
12201287
}
@@ -1261,6 +1328,12 @@ Result<std::shared_ptr<SparseTensor>> ReadSparseTensor(const Buffer& metadata,
12611328
type, shape, dim_names, checked_pointer_cast<SparseCSCIndex>(sparse_index),
12621329
non_zero_length, data);
12631330
}
1331+
case SparseTensorFormat::CSF: {
1332+
ARROW_ASSIGN_OR_RAISE(sparse_index, ReadSparseCSFIndex(sparse_tensor, shape, file));
1333+
return MakeSparseTensorWithSparseCSFIndex(
1334+
type, shape, dim_names, checked_pointer_cast<SparseCSFIndex>(sparse_index),
1335+
data);
1336+
}
12641337
default:
12651338
return Status::Invalid("Unsupported sparse index format");
12661339
}

‎cpp/src/arrow/ipc/writer.cc‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -771,6 +771,11 @@ class SparseTensorSerializer {
771771
VisitSparseCSCIndex(checked_cast<const SparseCSCIndex&>(sparse_index)));
772772
break;
773773

774+
case SparseTensorFormat::CSF:
775+
RETURN_NOT_OK(
776+
VisitSparseCSFIndex(checked_cast<const SparseCSFIndex&>(sparse_index)));
777+
break;
778+
774779
default:
775780
std::stringstream ss;
776781
ss << "Unable to convert type: " << sparse_index.ToString() << std::endl;
@@ -829,6 +834,16 @@ class SparseTensorSerializer {
829834
return Status::OK();
830835
}
831836

837+
Status VisitSparseCSFIndex(const SparseCSFIndex& sparse_index) {
838+
for (const std::shared_ptr<arrow::Tensor> indptr : sparse_index.indptr()) {
839+
out_->body_buffers.emplace_back(indptr->data());
840+
}
841+
for (const std::shared_ptr<arrow::Tensor> indices : sparse_index.indices()) {
842+
out_->body_buffers.emplace_back(indices->data());
843+
}
844+
return Status::OK();
845+
}
846+
832847
IpcPayload* out_;
833848

834849
std::vector<internal::BufferMetadata> buffer_meta_;

0 commit comments

Comments
 (0)