Add a new test of coo sparse tensor creation from non-contiguous tensor

Author: mrkn

cpp/src/arrow/sparse_tensor-test.cc

@@ -186,6 +186,58 @@ TEST(TestSparseCOOTensor, CreationFromTensor) {
   ASSERT_EQ(3, sidx->Value({11, 2}));
 }
 
+TEST(TestSparseCOOTensor, CreationFromNonContiguousTensor) {
+  std::vector<int64_t> shape = {2, 3, 4};
+  std::vector<int64_t> values = {1,  0, 0, 0, 2,  0, 0, 0, 0, 0, 3,  0, 0, 0, 4,  0,
+                                 5,  0, 0, 0, 6,  0, 0, 0, 0, 0, 11, 0, 0, 0, 12, 0,
+                                 13, 0, 0, 0, 14, 0, 0, 0, 0, 0, 15, 0, 0, 0, 16, 0};
+  std::vector<int64_t> strides = {192, 64, 16};
+  std::shared_ptr<Buffer> buffer = Buffer::Wrap(values);
+  Tensor tensor(int64(), buffer, shape, strides);
+  SparseTensorImpl<SparseCOOIndex> st(tensor);
+
+  ASSERT_EQ(12, st.non_zero_length());
+  ASSERT_TRUE(st.is_mutable());
+
+  const int64_t* ptr = reinterpret_cast<const int64_t*>(st.raw_data());
+  for (int i = 0; i < 6; ++i) {
+    ASSERT_EQ(i + 1, ptr[i]);
+  }
+  for (int i = 0; i < 6; ++i) {
+    ASSERT_EQ(i + 11, ptr[i + 6]);
+  }
+
+  const auto& si = internal::checked_cast<const SparseCOOIndex&>(*st.sparse_index());
+  std::shared_ptr<SparseCOOIndex::CoordsTensor> sidx = si.indices();
+  ASSERT_EQ(std::vector<int64_t>({12, 3}), sidx->shape());
+  ASSERT_TRUE(sidx->is_column_major());
+
+  // (0, 0, 0) -> 1
+  ASSERT_EQ(0, sidx->Value({0, 0}));
+  ASSERT_EQ(0, sidx->Value({0, 1}));
+  ASSERT_EQ(0, sidx->Value({0, 2}));
+
+  // (0, 0, 2) -> 2
+  ASSERT_EQ(0, sidx->Value({1, 0}));
+  ASSERT_EQ(0, sidx->Value({1, 1}));
+  ASSERT_EQ(2, sidx->Value({1, 2}));
+
+  // (0, 1, 1) -> 3
+  ASSERT_EQ(0, sidx->Value({2, 0}));
+  ASSERT_EQ(1, sidx->Value({2, 1}));
+  ASSERT_EQ(1, sidx->Value({2, 2}));
+
+  // (1, 2, 1) -> 15
+  ASSERT_EQ(1, sidx->Value({10, 0}));
+  ASSERT_EQ(2, sidx->Value({10, 1}));
+  ASSERT_EQ(1, sidx->Value({10, 2}));
+
+  // (1, 2, 3) -> 16
+  ASSERT_EQ(1, sidx->Value({11, 0}));
+  ASSERT_EQ(2, sidx->Value({11, 1}));
+  ASSERT_EQ(3, sidx->Value({11, 2}));
+}
+
 TEST(TestSparseCSRMatrix, CreationFromNumericTensor2D) {
   std::vector<int64_t> shape = {6, 4};
   std::vector<int64_t> values = {1, 0,  2, 0,  0,  3, 0,  4, 5, 0,  6, 0,