From fd4c81dca07d6953c07073d3e78cf2a54ea81048 Mon Sep 17 00:00:00 2001 From: Pratyush Adhikari Date: Wed, 19 Aug 2026 00:32:08 +0530 Subject: [PATCH] GH-50906: [C++] Support 1D tensors in SparseCSR/CSC matrix conversion --- cpp/src/arrow/sparse_tensor_test.cc | 64 +++++++++++++++++++++++++++ cpp/src/arrow/tensor/csx_converter.cc | 50 +++++++++++++++++---- 2 files changed, 106 insertions(+), 8 deletions(-) diff --git a/cpp/src/arrow/sparse_tensor_test.cc b/cpp/src/arrow/sparse_tensor_test.cc index 434f4a1723c7..0eaeedbbc22b 100644 --- a/cpp/src/arrow/sparse_tensor_test.cc +++ b/cpp/src/arrow/sparse_tensor_test.cc @@ -869,6 +869,70 @@ TEST_F(TestSparseCSRMatrix, TestToTensor) { ASSERT_TRUE(tensor.Equals(*dense_tensor)); } +TEST_F(TestSparseCSRMatrix, CreationFromNumericTensor1D) { + // Dense 1D vector: [1, 0, 2, 0, 0, 3, 0, 4] + // Expected: indptr = [0, 4], indices = [0, 2, 5, 7], values = [1, 2, 3, 4] + std::vector values = {1, 0, 2, 0, 0, 3, 0, 4}; + std::vector shape({static_cast(values.size())}); + NumericTensor dense_vector(Buffer::Wrap(values), shape); + + std::shared_ptr st; + ASSERT_OK_AND_ASSIGN(st, SparseCSRMatrix::Make(dense_vector)); + + ASSERT_EQ(4, st->non_zero_length()); + ASSERT_TRUE(st->is_mutable()); + + const int64_t* raw_data = reinterpret_cast(st->raw_data()); + AssertNumericDataEqual(raw_data, {1, 2, 3, 4}); + + auto si = internal::checked_pointer_cast(st->sparse_index()); + ASSERT_EQ(1, si->indptr()->ndim()); + ASSERT_EQ(1, si->indices()->ndim()); + + const int64_t* indptr_begin = + reinterpret_cast(si->indptr()->raw_data()); + std::vector indptr_values(indptr_begin, + indptr_begin + si->indptr()->shape()[0]); + ASSERT_EQ(std::vector({0, 4}), indptr_values); + + const int64_t* indices_begin = + reinterpret_cast(si->indices()->raw_data()); + std::vector indices_values(indices_begin, + indices_begin + si->indices()->shape()[0]); + ASSERT_EQ(std::vector({0, 2, 5, 7}), indices_values); +} + +TEST_F(TestSparseCSRMatrix, CreationFromAllZeroTensor1D) { + std::vector values = {0, 0, 0, 0}; + std::vector shape({static_cast(values.size())}); + NumericTensor dense_vector(Buffer::Wrap(values), shape); + + std::shared_ptr st; + ASSERT_OK_AND_ASSIGN(st, SparseCSRMatrix::Make(dense_vector)); + + ASSERT_EQ(0, st->non_zero_length()); + + auto si = internal::checked_pointer_cast(st->sparse_index()); + const int64_t* indptr_begin = + reinterpret_cast(si->indptr()->raw_data()); + std::vector indptr_values(indptr_begin, + indptr_begin + si->indptr()->shape()[0]); + ASSERT_EQ(std::vector({0, 0}), indptr_values); +} + +TEST_F(TestSparseCSRMatrix, TestToTensor1D) { + std::vector values = {1, 0, 2, 0, 0, 3, 0, 4}; + std::vector shape({static_cast(values.size())}); + Tensor tensor(int64(), Buffer::Wrap(values), shape); + + std::shared_ptr st; + ASSERT_OK_AND_ASSIGN(st, SparseCSRMatrix::Make(tensor)); + + ASSERT_EQ(4, st->non_zero_length()); + ASSERT_OK_AND_ASSIGN(auto dense_tensor, st->ToTensor()); + ASSERT_TRUE(tensor.Equals(*dense_tensor)); +} + template class TestSparseCSRMatrixEquality : public TestSparseTensorBase { public: diff --git a/cpp/src/arrow/tensor/csx_converter.cc b/cpp/src/arrow/tensor/csx_converter.cc index 679c3a0f1acd..63592b056bc8 100644 --- a/cpp/src/arrow/tensor/csx_converter.cc +++ b/cpp/src/arrow/tensor/csx_converter.cc @@ -56,13 +56,10 @@ class SparseCSXMatrixConverter : private SparseTensorConverterMixin { const int value_elsize = tensor_.type()->byte_width(); const int64_t ndim = tensor_.ndim(); - if (ndim > 2) { + if (ndim == 0 || ndim > 2) { return Status::Invalid("Invalid tensor dimension"); } - const int major_axis = static_cast(axis_); - const int64_t n_major = tensor_.shape()[major_axis]; - const int64_t n_minor = tensor_.shape()[1 - major_axis]; ARROW_ASSIGN_OR_RAISE(int64_t nonzero_count, tensor_.CountNonZero()); std::shared_ptr indptr_buffer; @@ -74,9 +71,45 @@ class SparseCSXMatrixConverter : private SparseTensorConverterMixin { const auto* tensor_data = tensor_.raw_data(); - if (ndim <= 1) { - return Status::NotImplemented("TODO for ndim <= 1"); + int64_t indptr_length; // n_major + 1 for 2D; 2 for 1D + + if (ndim == 1) { + // A 1D vector is treated as a single-row (CSR) or single-column (CSC) + // matrix. indptr has exactly 2 entries: [0, nnz]. + const int64_t n = tensor_.shape()[0]; + + ARROW_ASSIGN_OR_RAISE(indptr_buffer, AllocateBuffer(index_elsize * 2, pool_)); + auto* indptr = indptr_buffer->mutable_data(); + + ARROW_ASSIGN_OR_RAISE(indices_buffer, + AllocateBuffer(index_elsize * nonzero_count, pool_)); + auto* indices = indices_buffer->mutable_data(); + + // indptr[0] = 0 + AssignIndex(indptr, 0, index_elsize); + indptr += index_elsize; + + int64_t k = 0; + for (int64_t i = 0; i < n; ++i) { + const int64_t offset = i * value_elsize; + if (std::any_of(tensor_data + offset, tensor_data + offset + value_elsize, + IsNonZero)) { + std::copy_n(tensor_data + offset, value_elsize, values); + values += value_elsize; + AssignIndex(indices, i, index_elsize); + indices += index_elsize; + k++; + } + } + + // indptr[1] = nnz + AssignIndex(indptr, k, index_elsize); + indptr_length = 2; } else { + const int major_axis = static_cast(axis_); + const int64_t n_major = tensor_.shape()[major_axis]; + const int64_t n_minor = tensor_.shape()[1 - major_axis]; + ARROW_ASSIGN_OR_RAISE(indptr_buffer, AllocateBuffer(index_elsize * (n_major + 1), pool_)); auto* indptr = indptr_buffer->mutable_data(); @@ -111,9 +144,10 @@ class SparseCSXMatrixConverter : private SparseTensorConverterMixin { AssignIndex(indptr, k, index_elsize); indptr += index_elsize; } + indptr_length = n_major + 1; } - std::vector indptr_shape({n_major + 1}); + std::vector indptr_shape({indptr_length}); std::shared_ptr indptr_tensor = std::make_shared(index_value_type_, indptr_buffer, indptr_shape); @@ -181,7 +215,7 @@ Result> MakeTensorFromSparseCSXMatrix( std::vector strides; RETURN_NOT_OK(ComputeRowMajorStrides(fw_value_type, shape, &strides)); - const auto nc = shape[1]; + const auto nc = shape.size() > 1 ? shape[1] : shape[0]; int64_t offset = 0; for (int64_t i = 0; i < indptr->size() - 1; ++i) {