Skip to content

Commit

Permalink
ARROW-4226: [C++] Add sparse CSF tensor support
Browse files Browse the repository at this point in the history
This is to resolve [ARROW-4226](https://issues.apache.org/jira/browse/ARROW-4226).

Closes apache#5716 from rok/ARROW-4226 and squashes the following commits:

9ca93ab <Rok> Implementing review feedback.
1b922f6 <Rok> Implementing review feedback.
11b81bb <Rok> Factoring out index incrementing for dense to COO and CSF indices.
6f4f4a8 <Rok> Implementing feedback review.
28d38cb <Rok> Removing backslashes from comments.
3291abc <Rok> Marking indptrBuffers, indicesBuffers and axisOrder required.
d9ff47e <Rok> Further work and implementing review feedback.
24a831f <Rok> Style.
4f2bf00 <Rok> Work on CSF index tests.
6ceb406 <Rok> Implementing review feedback.
bd0d8c2 <Rok> Dense to sparse CSF conversion now in order of dimension size.
eb51947 <Rok> Switching SparseCSFIndex to '2D' data structure.
a322ff5 <Rok> Adding tests for multiple index value types for SparseCSFIndex.
f44d92c <Rok> Adding SparseCSFIndex::Make.
7d17995 <Rok> Adding Tensor to SparseCSFTensor conversion.
05a47a5 <Rok> Using axis_order in CSF.
6b938f7 <Rok> Documentation.
2d10104 <Rok> WIP

Authored-by: Rok <rok@mihevc.org>
Signed-off-by: Antoine Pitrou <antoine@python.org>
  • Loading branch information
rok authored and pitrou committed Feb 5, 2020
1 parent 25fd97b commit c02d376
Show file tree
Hide file tree
Showing 6 changed files with 628 additions and 25 deletions.
12 changes: 12 additions & 0 deletions cpp/src/arrow/compare.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1196,6 +1196,13 @@ inline bool SparseTensorEqualsImplDispatch(const SparseTensorImpl<SparseIndexTyp
right_csc);
}

case SparseTensorFormat::CSF: {
const auto& right_csf =
checked_cast<const SparseTensorImpl<SparseCSFIndex>&>(right);
return SparseTensorEqualsImpl<SparseIndexType, SparseCSFIndex>::Compare(left,
right_csf);
}

default:
return false;
}
Expand Down Expand Up @@ -1232,6 +1239,11 @@ bool SparseTensorEquals(const SparseTensor& left, const SparseTensor& right) {
return SparseTensorEqualsImplDispatch(left_csc, right);
}

case SparseTensorFormat::CSF: {
const auto& left_csf = checked_cast<const SparseTensorImpl<SparseCSFIndex>&>(left);
return SparseTensorEqualsImplDispatch(left_csf, right);
}

default:
return false;
}
Expand Down
5 changes: 5 additions & 0 deletions cpp/src/arrow/python/serialize.cc
Original file line number Diff line number Diff line change
Expand Up @@ -654,6 +654,7 @@ Status CountSparseTensors(
OwnedRef num_sparse_tensors(PyDict_New());
size_t num_coo = 0;
size_t num_csr = 0;
size_t num_csf = 0;

for (const auto& sparse_tensor : sparse_tensors) {
switch (sparse_tensor->format_id()) {
Expand All @@ -663,6 +664,9 @@ Status CountSparseTensors(
case SparseTensorFormat::CSR:
++num_csr;
break;
case SparseTensorFormat::CSF:
++num_csf;
break;
case SparseTensorFormat::CSC:
// TODO(mrkn): support csc
break;
Expand All @@ -671,6 +675,7 @@ Status CountSparseTensors(

PyDict_SetItemString(num_sparse_tensors.obj(), "coo", PyLong_FromSize_t(num_coo));
PyDict_SetItemString(num_sparse_tensors.obj(), "csr", PyLong_FromSize_t(num_csr));
PyDict_SetItemString(num_sparse_tensors.obj(), "csf", PyLong_FromSize_t(num_csf));
RETURN_IF_PYERROR();

*out = num_sparse_tensors.detach();
Expand Down
Loading

0 comments on commit c02d376

Please sign in to comment.