Implementing review feedback.

Author: rok

cpp/src/arrow/python/deserialize.cc

@@ -187,15 +187,15 @@ Status GetValue(PyObject* context, const Array& arr, int64_t index, int8_t type,
     case PythonType::SPARSECOOTENSOR: {
       int32_t ref = checked_cast<const Int32Array&>(arr).Value(index);
       const std::shared_ptr<SparseCOOTensor>& sparse_coo_tensor =
-          reinterpret_cast<const std::shared_ptr<SparseCOOTensor>&>(
+          arrow::internal::checked_pointer_cast<SparseCOOTensor>(
               blobs.sparse_tensors[ref]);
       *result = wrap_sparse_coo_tensor(sparse_coo_tensor);
       return Status::OK();
     }
     case PythonType::SPARSECSRMATRIX: {
       int32_t ref = checked_cast<const Int32Array&>(arr).Value(index);
       const std::shared_ptr<SparseCSRMatrix>& sparse_csr_matrix =
-          reinterpret_cast<const std::shared_ptr<SparseCSRMatrix>&>(
+          arrow::internal::checked_pointer_cast<SparseCSRMatrix>(
               blobs.sparse_tensors[ref]);
       *result = wrap_sparse_csr_matrix(sparse_csr_matrix);
       return Status::OK();
@@ -389,6 +389,7 @@ Status GetSerializedFromComponents(int num_tensors,
   }
 
   auto GetBuffer = [&data](Py_ssize_t index, std::shared_ptr<Buffer>* out) {
+    ARROW_CHECK_LE(index, PyList_Size(data));
     PyObject* py_buf = PyList_GET_ITEM(data, index);
     return unwrap_buffer(py_buf, out);
   };

cpp/src/arrow/python/deserialize.h

@@ -62,7 +62,7 @@ Status ReadSerializedObject(io::RandomAccessFile* src, SerializedPyObject* out);
 /// \param[in] num_ndarrays number of numpy Ndarrays in the object
 /// \param[in] num_buffers number of buffers in the object
 /// \param[in] data a list containing pyarrow.Buffer instances. Must be 1 +
-/// num_tensors * 2 + num_sparse_tensors * 2 + num_buffers in length
+/// num_tensors * 2 + num_coo_tensors * 3 + num_csr_tensors * 4 + num_buffers in length
 /// \param[out] out the reconstructed object
 /// \return Status
 ARROW_PYTHON_EXPORT

cpp/src/arrow/python/serialize.h

@@ -140,9 +140,9 @@ struct PythonType {
     TENSOR,
     NDARRAY,
     BUFFER,
-    NUM_PYTHON_TYPES,
     SPARSECOOTENSOR,
-    SPARSECSRMATRIX
+    SPARSECSRMATRIX,
+    NUM_PYTHON_TYPES
   };
 };
 

python/pyarrow/tests/test_serialization.py

@@ -152,11 +152,9 @@ def assert_equal(obj1, obj2):
 ]
 
 
-# index_types = ('i1', 'i2', 'i4', 'i8', 'u1', 'u2', 'u4', 'u8')
-# tensor_types = ('i1', 'i2', 'i4', 'i8', 'u1', 'u2', 'u4', 'u8',
-#                 'f2', 'f4', 'f8')
-index_types = ('i8',)
-tensor_types = ('i8',)
+index_types = ('i1', 'i2', 'i4', 'i8', 'u1', 'u2', 'u4', 'u8')
+tensor_types = ('i1', 'i2', 'i4', 'i8', 'u1', 'u2', 'u4', 'u8',
+                'f2', 'f4', 'f8')
 
 
 if sys.version_info >= (3, 0):
@@ -544,10 +542,10 @@ def deserializer(data):
 def test_sparse_coo_tensor_serialization(index_type, tensor_type):
     tensor_dtype = np.dtype(tensor_type)
     index_dtype = np.dtype(index_type)
-    data = np.array([[1, 2, 3, 4, 5, 6, 7]]).T.astype(tensor_dtype)
+    data = np.array([[1, 2, 3, 4, 5, 6]]).T.astype(tensor_dtype)
     coords = np.array([
-        [0, 0, 2, 3, 1, 3, 0],
-        [0, 2, 0, 4, 5, 5, 0],
+        [0, 0, 2, 3, 1, 3],
+        [0, 2, 0, 4, 5, 5],
     ]).T.astype(index_dtype)
     shape = (4, 6)
     dim_names = ('x', 'y')
@@ -558,7 +556,7 @@ def test_sparse_coo_tensor_serialization(index_type, tensor_type):
     context = pa.default_serialization_context()
     serialized = pa.serialize(sparse_tensor, context=context).to_buffer()
     result = pa.deserialize(serialized)
-    result.equals(sparse_tensor)
+    assert_equal(result, sparse_tensor)
 
     data_result, coords_result = result.to_numpy()
     assert np.array_equal(data_result, data)
@@ -572,10 +570,10 @@ def test_sparse_coo_tensor_components_serialization(large_buffer,
                                                     index_type, tensor_type):
     tensor_dtype = np.dtype(tensor_type)
     index_dtype = np.dtype(index_type)
-    data = np.array([[1, 2, 3, 4, 5, 6, 7]]).T.astype(tensor_dtype)
+    data = np.array([[1, 2, 3, 4, 5, 6]]).T.astype(tensor_dtype)
     coords = np.array([
-        [0, 0, 2, 3, 1, 3, 0],
-        [0, 2, 0, 4, 5, 5, 0],
+        [0, 0, 2, 3, 1, 3],
+        [0, 2, 0, 4, 5, 5],
     ]).T.astype(index_dtype)
     shape = (4, 6)
     dim_names = ('x', 'y')
@@ -587,9 +585,9 @@ def test_sparse_coo_tensor_components_serialization(large_buffer,
 
 @pytest.mark.skipif(not coo_matrix, reason="requires scipy")
 def test_scipy_sparse_coo_tensor_serialization():
-    data = np.array([1, 2, 3, 4, 5, 6, 7])
-    row = np.array([0, 0, 2, 3, 1, 3, 0])
-    col = np.array([0, 2, 0, 4, 5, 5, 0])
+    data = np.array([1, 2, 3, 4, 5, 6])
+    row = np.array([0, 0, 2, 3, 1, 3])
+    col = np.array([0, 2, 0, 4, 5, 5])
     shape = (4, 6)
 
     sparse_array = coo_matrix((data, (row, col)), shape=shape)

python/pyarrow/tests/test_sparse_tensor.py

@@ -206,10 +206,10 @@ def test_sparse_csr_matrix_from_dense(dtype_str, arrow_type):
 @pytest.mark.parametrize('dtype_str,arrow_type', tensor_type_pairs)
 def test_sparse_coo_tensor_numpy_roundtrip(dtype_str, arrow_type):
     dtype = np.dtype(dtype_str)
-    data = np.array([[1, 2, 3, 4, 5, 6, 7]]).T.astype(dtype)
+    data = np.array([[1, 2, 3, 4, 5, 6]]).T.astype(dtype)
     coords = np.array([
-        [0, 0, 2, 3, 1, 3, 0],
-        [0, 2, 0, 4, 5, 5, 0],
+        [0, 0, 2, 3, 1, 3],
+        [0, 2, 0, 4, 5, 5],
     ]).T
     shape = (4, 6)
     dim_names = ('x', 'y')
@@ -271,9 +271,9 @@ def test_dense_to_sparse_tensor(dtype_str, arrow_type, sparse_tensor_type):
 @pytest.mark.parametrize('dtype_str,arrow_type', tensor_type_pairs)
 def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type):
     dtype = np.dtype(dtype_str)
-    data = np.array([1, 2, 3, 4, 5, 6, 7]).astype(dtype)
-    row = np.array([0, 0, 2, 3, 1, 3, 0])
-    col = np.array([0, 2, 0, 4, 5, 5, 0])
+    data = np.array([1, 2, 3, 4, 5, 6]).astype(dtype)
+    row = np.array([0, 0, 2, 3, 1, 3])
+    col = np.array([0, 2, 0, 4, 5, 5])
     shape = (4, 6)
     dim_names = ('x', 'y')
 
@@ -288,6 +288,10 @@ def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type):
     assert np.array_equal(sparse_array.data, out_sparse_array.data)
     assert np.array_equal(sparse_array.row, out_sparse_array.row)
     assert np.array_equal(sparse_array.col, out_sparse_array.col)
+    if dtype_str != 'f2':
+        assert np.array_equal(
+            sparse_array.astype(np.float32).toarray().astype(np.float16),
+            sparse_tensor.to_tensor().to_numpy())
 
 
 @pytest.mark.skipif(not csr_matrix, reason="requires scipy")
@@ -311,3 +315,7 @@ def test_sparse_csr_matrix_scipy_roundtrip(dtype_str, arrow_type):
     assert np.array_equal(sparse_array.data, out_sparse_array.data)
     assert np.array_equal(sparse_array.indptr, out_sparse_array.indptr)
     assert np.array_equal(sparse_array.indices, out_sparse_array.indices)
+    if dtype_str != 'f2':
+        assert np.array_equal(
+            sparse_array.astype(np.float32).toarray().astype(np.float16),
+            sparse_tensor.to_tensor().to_numpy())